version 0.0.17 (#111)

Release 0.0.17

README.md

@@ -1,4 +1,4 @@
-# LangServe 🦜️🔗
+# LangServe 🦜️🏓
 
 ## Overview
 
@@ -25,6 +25,10 @@ A javascript client is available in [LangChainJS](https://js.langchain.com/docs/
 - Client callbacks are not yet supported for events that originate on the server
 - Does not work with [pydantic v2 yet](https://github.com/tiangolo/fastapi/issues/10360)
 
+## Security
+
+* Vulnerability in Versions 0.0.13 - 0.0.15 -- playground endpoint allows accessing arbitrary files on server. [Resolved in 0.0.16](https://github.com/langchain-ai/langserve/pull/98).
+
 ## LangChain CLI 🛠️
 
 Use the `LangChain` CLI to bootstrap a `LangServe` project quickly.

langserve/callbacks.py

@@ -0,0 +1,475 @@
+from __future__ import annotations
+
+import uuid
+from typing import Any, Dict, List, Optional, Sequence
+from uuid import UUID
+
+from langchain.callbacks.base import AsyncCallbackHandler
+from langchain.callbacks.manager import (
+    BaseRunManager,
+    ahandle_event,
+    handle_event,
+)
+from langchain.schema import AgentAction, AgentFinish, BaseMessage, Document, LLMResult
+from typing_extensions import TypedDict
+
+
+class CallbackEventDict(TypedDict, total=False):
+    """A dictionary representation of a callback event."""
+
+    type: str
+    parent_run_id: Optional[UUID]
+    run_id: UUID
+
+
+class AsyncEventAggregatorCallback(AsyncCallbackHandler):
+    """A callback handler that aggregates all the events that have been called.
+
+    This callback handler aggregates all the events that have been called placing
+    them in a single mutable list.
+
+    This callback handler is not threading safe, and is meant to be used in an async
+    context only.
+    """
+
+    def __init__(self) -> None:
+        """Get a list of all the callback events that have been called."""
+        super().__init__()
+        # Callback events is a mutable state that is used only in an async context,
+        # so it should be safe to mutate without the usage of a lock.
+        self.callback_events: List[CallbackEventDict] = []
+
+    def log_callback(self, event: CallbackEventDict) -> None:
+        """Log the callback event."""
+        self.callback_events.append(event)
+
+    async def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Attempt to serialize the callback event."""
+        self.log_callback(
+            {
+                "type": "on_chat_model_start",
+                "serialized": serialized,
+                "messages": messages,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "metadata": metadata,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Attempt to serialize the callback event."""
+        self.log_callback(
+            {
+                "type": "on_chain_start",
+                "serialized": serialized,
+                "inputs": inputs,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "metadata": metadata,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_chain_end(
+        self,
+        outputs: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_chain_end",
+                "outputs": outputs,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_chain_error",
+                "error": error,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_retriever_start(
+        self,
+        serialized: Dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_retriever_start",
+                "serialized": serialized,
+                "query": query,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "metadata": metadata,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_retriever_end(
+        self,
+        documents: Sequence[Document],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_retriever_end",
+                "documents": documents,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_retriever_error",
+                "error": error,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_tool_start",
+                "serialized": serialized,
+                "input_str": input_str,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "metadata": metadata,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_tool_end",
+                "output": output,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_tool_error",
+                "error": error,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_agent_action",
+                "action": action,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_agent_finish",
+                "finish": finish,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_llm_start",
+                "serialized": serialized,
+                "prompts": prompts,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "metadata": metadata,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_llm_end",
+                "response": response,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+    async def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_callback(
+            {
+                "type": "on_llm_error",
+                "error": error,
+                "run_id": run_id,
+                "parent_run_id": parent_run_id,
+                "tags": tags,
+                "kwargs": kwargs,
+            }
+        )
+
+
+def replace_uuids(
+    callback_events: Sequence[CallbackEventDict],
+) -> List[CallbackEventDict]:
+    """Replace uids in the event callbacks with new uids.
+
+    This function mutates the event callback events in place.
+
+    Args:
+        callback_events: A list of event callbacks.
+    """
+    # Create a dictionary to store mappings from old UID to new UID
+    uid_mapping: dict = {}
+
+    updated_events = []
+
+    # Iterate through the list of event callbacks
+    for event in callback_events:
+        updated_event = event.copy()
+        # Replace UIDs in the 'run_id' field
+        if "run_id" in updated_event and updated_event["run_id"] is not None:
+            if updated_event["run_id"] not in uid_mapping:
+                # Generate a new UUID
+                new_uid = uuid.uuid4()
+                uid_mapping[updated_event["run_id"]] = new_uid
+            # Replace the old UID with the new one
+            updated_event["run_id"] = uid_mapping[updated_event["run_id"]]
+
+        # Replace UIDs in the 'parent_run_id' field if it's not None
+        if (
+            "parent_run_id" in updated_event
+            and updated_event["parent_run_id"] is not None
+        ):
+            if updated_event["parent_run_id"] not in uid_mapping:
+                # Generate a new UUID
+                new_uid = uuid.uuid4()
+                uid_mapping[updated_event["parent_run_id"]] = new_uid
+            # Replace the old UID with the new one
+            updated_event["parent_run_id"] = uid_mapping[updated_event["parent_run_id"]]
+        updated_events.append(updated_event)
+    return updated_events
+
+
+# Mapping from event name to ignore condition name
+NAME_TO_IGNORE_CONDITION = {
+    "on_retry": "ignore_retry",
+    "on_text": None,
+    "on_agent_action": "ignore_agent",
+    "on_agent_finish": "ignore_agent",
+    "on_llm_start": "ignore_llm",
+    "on_llm_end": "ignore_llm",
+    "on_llm_error": "ignore_llm",
+    "on_chain_start": "ignore_chain",
+    "on_chain_end": "ignore_chain",
+    "on_chain_error": "ignore_chain",
+    "on_chat_model_start": "ignore_chat_model",
+    "on_tool_start": "ignore_agent",
+    "on_tool_end": "ignore_agent",
+    "on_tool_error": "ignore_agent",
+    "on_retriever_start": "ignore_retriever",
+    "on_retriever_end": "ignore_retriever",
+    "on_retriever_error": "ignore_retriever",
+}
+
+
+async def ahandle_callbacks(
+    callback_manager: BaseRunManager,
+    callback_events: Sequence[CallbackEventDict],
+) -> None:
+    """Invoke all the callback handlers with the given callback events."""
+    callback_events = replace_uuids(callback_events)
+
+    # 1. Do I need inheritable handlers
+    for event in callback_events:
+        if event["parent_run_id"] is None:  # How do we make sure it's None!?
+            event["parent_run_id"] = callback_manager.run_id
+
+        event_data = {key: value for key, value in event.items() if key != "type"}
+
+        await ahandle_event(
+            # Unpacking like this may not work
+            callback_manager.handlers,
+            event["type"],
+            ignore_condition_name=NAME_TO_IGNORE_CONDITION.get(event["type"], None),
+            **event_data,
+        )
+
+
+def handle_callbacks(
+    callback_manager: BaseRunManager,
+    callback_events: Sequence[CallbackEventDict],
+) -> None:
+    """Invoke all the callback handlers with the given callback events."""
+    callback_events = replace_uuids(callback_events)
+
+    for event in callback_events:
+        if event["parent_run_id"] is None:  # How do we make sure it's None!?
+            event["parent_run_id"] = callback_manager.run_id
+
+        event_data = {key: value for key, value in event.items() if key != "type"}
+
+        handle_event(
+            # Unpacking like this may not work
+            callback_manager.handlers,
+            event["type"],
+            ignore_condition_name=NAME_TO_IGNORE_CONDITION.get(event["type"], None),
+            **event_data,
+        )

langserve/client.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import copy
 import json
 import logging
 import weakref
@@ -14,12 +15,17 @@ from typing import (
     List,
     Optional,
     Sequence,
+    Tuple,
     Union,
 )
 from urllib.parse import urljoin
 
 import httpx
 from httpx._types import AuthTypes, CertTypes, CookieTypes, HeaderTypes, VerifyTypes
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+)
 from langchain.callbacks.tracers.log_stream import RunLogPatch
 from langchain.load.dump import dumpd
 from langchain.schema.runnable import Runnable
@@ -31,7 +37,12 @@ from langchain.schema.runnable.config import (
 )
 from langchain.schema.runnable.utils import Input, Output
 
-from langserve.serialization import simple_dumpd, simple_loads
+from langserve.callbacks import CallbackEventDict, ahandle_callbacks, handle_callbacks
+from langserve.serialization import (
+    Serializer,
+    WellKnownLCSerializer,
+    load_events,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -42,12 +53,35 @@ def _without_callbacks(config: Optional[RunnableConfig]) -> RunnableConfig:
     return {k: v for k, v in _config.items() if k != "callbacks"}
 
 
-@lru_cache(maxsize=1_000)  # Will accommodate up to 100 different error messages
+@lru_cache(maxsize=1_000)  # Will accommodate up to 1_000 different error messages
 def _log_error_message_once(error_message: str) -> None:
     """Log an error message once."""
     logger.error(error_message)
 
 
+def _sanitize_request(request: httpx.Request) -> httpx.Request:
+    """Remove sensitive headers from the request."""
+    accept_headers = {
+        "accept",
+        "content-type",
+        "user-agent",
+        "connection",
+        "content-length",
+        "accept-encoding",
+        "host",
+    }
+    new_headers = request.headers.copy()
+    for key, value in new_headers.items():
+        if key.lower() not in accept_headers:
+            new_headers[key] = "<redacted>"
+        else:
+            new_headers[key] = value
+
+    new_request = copy.copy(request)
+    new_request.headers = new_headers
+    return new_request
+
+
 def _raise_for_status(response: httpx.Response) -> None:
     """Re-raise with a more informative message.
 
@@ -69,7 +103,7 @@ def _raise_for_status(response: httpx.Response) -> None:
 
         raise httpx.HTTPStatusError(
             message=message,
-            request=e.request,
+            request=_sanitize_request(e.request),
             response=e.response,
         )
 
@@ -112,12 +146,12 @@ def _raise_exception_from_data(data: str, request: httpx.Request) -> None:
     except json.JSONDecodeError:
         raise httpx.HTTPStatusError(
             message="invalid json in error event sent from server",
-            request=request,
+            request=_sanitize_request(request),
             response=httpx.Response(status_code=500, text=data),
         )
     raise httpx.HTTPStatusError(
         message=decoded_data["message"],
-        request=request,
+        request=_sanitize_request(request),
         response=httpx.Response(
             status_code=decoded_data["status_code"],
             text=decoded_data["message"],
@@ -125,6 +159,42 @@ def _raise_exception_from_data(data: str, request: httpx.Request) -> None:
     )
 
 
+def _decode_response(
+    serializer: Serializer,
+    response: httpx.Response,
+    *,
+    is_batch: bool = False,
+) -> Tuple[Any, Union[List[CallbackEventDict], List[List[CallbackEventDict]]]]:
+    """Decode the response."""
+    _raise_for_status(response)
+    obj = response.json()
+    if not isinstance(obj, dict):
+        raise ValueError(f"Expected a dictionary, got {obj}")
+
+    if "output" not in obj:
+        raise ValueError("Key `output` not found in")
+
+    output = serializer.loadd(obj["output"])
+
+    if "callback_events" in obj:
+        if is_batch:
+            if not isinstance(obj["callback_events"], list):
+                raise ValueError(
+                    f"Expected a list of callback events, got {obj['callback_events']}"
+                )
+            else:
+                callback_events = [
+                    load_events(callback_events)
+                    for callback_events in obj["callback_events"]
+                ]
+        else:
+            callback_events = load_events(obj["callback_events"])
+    else:
+        callback_events = []
+
+    return output, callback_events
+
+
 class RemoteRunnable(Runnable[Input, Output]):
     """A RemoteRunnable is a runnable that is executed on a remote server.
 
@@ -134,8 +204,6 @@ class RemoteRunnable(Runnable[Input, Output]):
 
     - `batch` with `return_exceptions=True` since we do not support exception
       translation from the server.
-    - Callbacks via the `config` argument as serialization of callbacks is not
-      supported.
     """
 
     def __init__(
@@ -149,6 +217,7 @@ class RemoteRunnable(Runnable[Input, Output]):
         verify: VerifyTypes = True,
         cert: Optional[CertTypes] = None,
         client_kwargs: Optional[Dict[str, Any]] = None,
+        use_server_callback_events: bool = True,
     ) -> None:
         """Initialize the client.
 
@@ -161,7 +230,9 @@ class RemoteRunnable(Runnable[Input, Output]):
             verify: Whether to verify SSL certificates
             cert: SSL certificate to use for requests
             client_kwargs: If provided will be unpacked as kwargs to both the sync
-                           and async httpx clients
+                and async httpx clients
+            use_server_callback_events: Whether to invoke callbacks on any
+                callback events returned by the server.
         """
         _client_kwargs = client_kwargs or {}
         self.url = url
@@ -188,21 +259,32 @@ class RemoteRunnable(Runnable[Input, Output]):
 
         # Register cleanup handler once RemoteRunnable is garbage collected
         weakref.finalize(self, _close_clients, self.sync_client, self.async_client)
+        self._lc_serializer = WellKnownLCSerializer()
+        self._use_server_callback_events = use_server_callback_events
 
     def _invoke(
-        self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
+        self,
+        input: Input,
+        run_manager: CallbackManagerForChainRun,
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
     ) -> Output:
         """Invoke the runnable with the given input and config."""
         response = self.sync_client.post(
             "/invoke",
             json={
-                "input": simple_dumpd(input),
+                "input": self._lc_serializer.dumpd(input),
                 "config": _without_callbacks(config),
                 "kwargs": kwargs,
             },
         )
-        _raise_for_status(response)
-        return simple_loads(response.text)["output"]
+        output, callback_events = _decode_response(
+            self._lc_serializer, response, is_batch=False
+        )
+
+        if self._use_server_callback_events and callback_events:
+            handle_callbacks(run_manager, callback_events)
+        return output
 
     def invoke(
         self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -212,18 +294,27 @@ class RemoteRunnable(Runnable[Input, Output]):
         return self._call_with_config(self._invoke, input, config=config)
 
     async def _ainvoke(
-        self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
+        self,
+        input: Input,
+        run_manager: AsyncCallbackManagerForChainRun,
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
     ) -> Output:
+        """Invoke the runnable with the given input and config."""
         response = await self.async_client.post(
             "/invoke",
             json={
-                "input": simple_dumpd(input),
+                "input": self._lc_serializer.dumpd(input),
                 "config": _without_callbacks(config),
                 "kwargs": kwargs,
             },
         )
-        _raise_for_status(response)
-        return simple_loads(response.text)["output"]
+        output, callback_events = _decode_response(
+            self._lc_serializer, response, is_batch=False
+        )
+        if self._use_server_callback_events and callback_events:
+            handle_callbacks(run_manager, callback_events)
+        return output
 
     async def ainvoke(
         self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -235,6 +326,7 @@ class RemoteRunnable(Runnable[Input, Output]):
     def _batch(
         self,
         inputs: List[Input],
+        run_manager: List[CallbackManagerForChainRun],
         config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
         *,
         return_exceptions: bool = False,
@@ -255,13 +347,23 @@ class RemoteRunnable(Runnable[Input, Output]):
         response = self.sync_client.post(
             "/batch",
             json={
-                "inputs": simple_dumpd(inputs),
+                "inputs": self._lc_serializer.dumpd(inputs),
                 "config": _config,
                 "kwargs": kwargs,
             },
         )
-        _raise_for_status(response)
-        return simple_loads(response.text)["output"]
+        outputs, corresponding_callback_events = _decode_response(
+            self._lc_serializer, response, is_batch=True
+        )
+
+        # Now handle callbacks if any were returned
+        if self._use_server_callback_events and corresponding_callback_events:
+            for run_manager_, callback_events in zip(
+                run_manager, corresponding_callback_events
+            ):
+                handle_callbacks(run_manager_, callback_events)
+
+        return outputs
 
     def batch(
         self,
@@ -276,6 +378,7 @@ class RemoteRunnable(Runnable[Input, Output]):
     async def _abatch(
         self,
         inputs: List[Input],
+        run_manager: List[AsyncCallbackManagerForChainRun],
         config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
         *,
         return_exceptions: bool = False,
@@ -297,13 +400,27 @@ class RemoteRunnable(Runnable[Input, Output]):
         response = await self.async_client.post(
             "/batch",
             json={
-                "inputs": simple_dumpd(inputs),
+                "inputs": self._lc_serializer.dumpd(inputs),
                 "config": _config,
                 "kwargs": kwargs,
             },
         )
-        _raise_for_status(response)
-        return simple_loads(response.text)["output"]
+        outputs, corresponding_callback_events = _decode_response(
+            self._lc_serializer, response, is_batch=True
+        )
+
+        # Now handle callbacks
+
+        if self._use_server_callback_events and corresponding_callback_events:
+            tasks = []
+            for run_manager_, callback_events in zip(
+                run_manager, corresponding_callback_events
+            ):
+                tasks.append(ahandle_callbacks(run_manager_, callback_events))
+
+            # Execute coroutines concurrently
+            await asyncio.gather(*tasks)
+        return outputs
 
     async def abatch(
         self,
@@ -334,11 +451,11 @@ class RemoteRunnable(Runnable[Input, Output]):
 
         run_manager = callback_manager.on_chain_start(
             dumpd(self),
-            simple_dumpd(input),
+            self._lc_serializer.dumpd(input),
             name=config.get("run_name"),
         )
         data = {
-            "input": simple_dumpd(input),
+            "input": self._lc_serializer.dumpd(input),
             "config": _without_callbacks(config),
             "kwargs": kwargs,
         }
@@ -358,7 +475,7 @@ class RemoteRunnable(Runnable[Input, Output]):
             ) as event_source:
                 for sse in event_source.iter_sse():
                     if sse.event == "data":
-                        chunk = simple_loads(sse.data)
+                        chunk = self._lc_serializer.loads(sse.data)
                         yield chunk
 
                         if final_output:
@@ -397,11 +514,11 @@ class RemoteRunnable(Runnable[Input, Output]):
 
         run_manager = await callback_manager.on_chain_start(
             dumpd(self),
-            simple_dumpd(input),
+            self._lc_serializer.dumpd(input),
             name=config.get("run_name"),
         )
         data = {
-            "input": simple_dumpd(input),
+            "input": self._lc_serializer.dumpd(input),
             "config": _without_callbacks(config),
             "kwargs": kwargs,
         }
@@ -418,7 +535,7 @@ class RemoteRunnable(Runnable[Input, Output]):
             ) as event_source:
                 async for sse in event_source.aiter_sse():
                     if sse.event == "data":
-                        chunk = simple_loads(sse.data)
+                        chunk = self._lc_serializer.loads(sse.data)
                         yield chunk
 
                         if final_output:
@@ -477,11 +594,11 @@ class RemoteRunnable(Runnable[Input, Output]):
 
         run_manager = await callback_manager.on_chain_start(
             dumpd(self),
-            simple_dumpd(input),
+            self._lc_serializer.dumpd(input),
             name=config.get("run_name"),
         )
         data = {
-            "input": simple_dumpd(input),
+            "input": self._lc_serializer.dumpd(input),
             "config": _without_callbacks(config),
             "kwargs": kwargs,
             "diff": True,
@@ -505,7 +622,7 @@ class RemoteRunnable(Runnable[Input, Output]):
             ) as event_source:
                 async for sse in event_source.aiter_sse():
                     if sse.event == "data":
-                        data = simple_loads(sse.data)
+                        data = self._lc_serializer.loads(sse.data)
                         chunk = RunLogPatch(*data["ops"])
 
                         yield chunk

langserve/playground.py

@@ -2,7 +2,7 @@ import json
 import mimetypes
 import os
 from string import Template
-from typing import List, Type
+from typing import Sequence, Type
 
 from fastapi.responses import Response
 from langchain.schema.runnable import Runnable
@@ -20,28 +20,42 @@ class PlaygroundTemplate(Template):
 async def serve_playground(
     runnable: Runnable,
     input_schema: Type[BaseModel],
-    config_keys: List[str],
+    config_keys: Sequence[str],
     base_url: str,
     file_path: str,
 ) -> Response:
-    local_file_path = os.path.join(
-        os.path.dirname(__file__),
-        "./playground/dist",
-        file_path or "index.html",
+    """Serve the playground."""
+    local_file_path = os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "./playground/dist",
+            file_path or "index.html",
+        )
     )
-    with open(local_file_path) as f:
-        mime_type = mimetypes.guess_type(local_file_path)[0]
-        if mime_type in ("text/html", "text/css", "application/javascript"):
-            res = PlaygroundTemplate(f.read()).substitute(
-                LANGSERVE_BASE_URL=base_url[1:]
-                if base_url.startswith("/")
-                else base_url,
-                LANGSERVE_CONFIG_SCHEMA=json.dumps(
-                    runnable.config_schema(include=config_keys).schema()
-                ),
-                LANGSERVE_INPUT_SCHEMA=json.dumps(input_schema.schema()),
-            )
-        else:
-            res = f.buffer.read()
 
-    return Response(res, media_type=mime_type)
+    base_dir = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "./playground/dist")
+    )
+
+    if base_dir != os.path.commonpath((base_dir, local_file_path)):
+        return Response("Not Found", status_code=404)
+
+    try:
+        with open(local_file_path) as f:
+            mime_type = mimetypes.guess_type(local_file_path)[0]
+            if mime_type in ("text/html", "text/css", "application/javascript"):
+                response = PlaygroundTemplate(f.read()).substitute(
+                    LANGSERVE_BASE_URL=base_url[1:]
+                    if base_url.startswith("/")
+                    else base_url,
+                    LANGSERVE_CONFIG_SCHEMA=json.dumps(
+                        runnable.config_schema(include=config_keys).schema()
+                    ),
+                    LANGSERVE_INPUT_SCHEMA=json.dumps(input_schema.schema()),
+                )
+            else:
+                response = f.buffer.read()
+    except FileNotFoundError:
+        return Response("Not Found", status_code=404)
+
+    return Response(response, media_type=mime_type)

langserve/playground/.gitignore

@@ -23,4 +23,6 @@ dist-ssr
 *.sln
 *.sw?
 
-.yarn
+.yarn
+
+!dist

langserve/playground/dist/index.html

@@ -5,8 +5,8 @@
     <link rel="icon" href="/____LANGSERVE_BASE_URL/favicon.ico" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Playground</title>
-    <script type="module" crossorigin src="/____LANGSERVE_BASE_URL/assets/index-9fe7f71c.js"></script>
-    <link rel="stylesheet" href="/____LANGSERVE_BASE_URL/assets/index-5008c8a8.css">
+    <script type="module" crossorigin src="/____LANGSERVE_BASE_URL/assets/index-0ea7a041.js"></script>
+    <link rel="stylesheet" href="/____LANGSERVE_BASE_URL/assets/index-da983f33.css">
   </head>
   <body>
     <div id="root"></div>

langserve/playground/src/App.tsx

@@ -42,7 +42,6 @@ import {
   vanillaRenderers,
   InputControl,
 } from "@jsonforms/vanilla-renderers";
-
 import { useSchemas } from "./useSchemas";
 import { RunState, useStreamLog } from "./useStreamLog";
 import {
@@ -61,6 +60,14 @@ import CustomTextAreaCell from "./components/CustomTextAreaCell";
 import JsonTextAreaCell from "./components/JsonTextAreaCell";
 import { cn } from "./utils/cn";
 import { getStateFromUrl, ShareDialog } from "./components/ShareDialog";
+import {
+  chatMessagesTester,
+  ChatMessagesControlRenderer,
+} from "./components/ChatMessagesControlRenderer";
+import {
+  ChatMessageTuplesControlRenderer,
+  chatMessagesTupleTester,
+} from "./components/ChatMessageTuplesControlRenderer";
 
 dayjs.extend(relativeDate);
 dayjs.extend(utc);
@@ -98,12 +105,119 @@ const renderers = [
   // custom renderers
   { tester: materialArrayControlTester, renderer: CustomArrayControlRenderer },
   { tester: isObject, renderer: InputControl },
+  { tester: chatMessagesTester, renderer: ChatMessagesControlRenderer },
+  {
+    tester: chatMessagesTupleTester,
+    renderer: ChatMessageTuplesControlRenderer,
+  },
 ];
 
 const nestedArrayControlTester: RankedTester = rankWith(1, (_, jsonSchema) => {
   return jsonSchema.type === "array";
 });
 
+// inlined from langchain/schema
+interface BaseMessageFields {
+  content: string;
+  name?: string;
+  additional_kwargs?: {
+    [key: string]: unknown;
+  };
+}
+
+class AIMessageChunk {
+  /** The text of the message. */
+  content: string;
+
+  /** The name of the message sender in a multi-user chat. */
+  name?: string;
+
+  /** Additional keyword arguments */
+  additional_kwargs: NonNullable<BaseMessageFields["additional_kwargs"]>;
+
+  constructor(fields: BaseMessageFields) {
+    // Make sure the default value for additional_kwargs is passed into super() for serialization
+    if (!fields.additional_kwargs) {
+      // eslint-disable-next-line no-param-reassign
+      fields.additional_kwargs = {};
+    }
+
+    this.name = fields.name;
+    this.content = fields.content;
+    this.additional_kwargs = fields.additional_kwargs;
+  }
+
+  static _mergeAdditionalKwargs(
+    left: NonNullable<BaseMessageFields["additional_kwargs"]>,
+    right: NonNullable<BaseMessageFields["additional_kwargs"]>
+  ): NonNullable<BaseMessageFields["additional_kwargs"]> {
+    const merged = { ...left };
+    for (const [key, value] of Object.entries(right)) {
+      if (merged[key] === undefined) {
+        merged[key] = value;
+      } else if (typeof merged[key] !== typeof value) {
+        throw new Error(
+          `additional_kwargs[${key}] already exists in the message chunk, but with a different type.`
+        );
+      } else if (typeof merged[key] === "string") {
+        merged[key] = (merged[key] as string) + value;
+      } else if (
+        !Array.isArray(merged[key]) &&
+        typeof merged[key] === "object"
+      ) {
+        merged[key] = this._mergeAdditionalKwargs(
+          merged[key] as NonNullable<BaseMessageFields["additional_kwargs"]>,
+          value as NonNullable<BaseMessageFields["additional_kwargs"]>
+        );
+      } else {
+        throw new Error(
+          `additional_kwargs[${key}] already exists in this message chunk.`
+        );
+      }
+    }
+    return merged;
+  }
+
+  concat(chunk: AIMessageChunk) {
+    return new AIMessageChunk({
+      content: this.content + chunk.content,
+      additional_kwargs: AIMessageChunk._mergeAdditionalKwargs(
+        this.additional_kwargs,
+        chunk.additional_kwargs
+      ),
+    });
+  }
+}
+
+function isAiMessageChunkFields(value: unknown): value is BaseMessageFields {
+  if (typeof value !== "object" || value == null) return false;
+  return "content" in value && typeof value["content"] === "string";
+}
+
+function isAiMessageChunkFieldsList(
+  value: unknown[]
+): value is BaseMessageFields[] {
+  return value.length > 0 && value.every((x) => isAiMessageChunkFields(x));
+}
+
+function StreamOutput(props: { streamed: unknown[] }) {
+  // check if we're streaming AIMessageChunk
+  if (isAiMessageChunkFieldsList(props.streamed)) {
+    const concat = props.streamed.reduce<AIMessageChunk | null>(
+      (memo, field) => {
+        const chunk = new AIMessageChunk(field);
+        if (memo == null) return chunk;
+        return memo.concat(chunk);
+      },
+      null
+    );
+
+    return concat?.content || "...";
+  }
+
+  return props.streamed.map(str).join("") || "...";
+}
+
 const cells = [
   { tester: booleanCellTester, cell: BooleanCell },
   { tester: dateCellTester, cell: DateCell },
@@ -184,7 +298,7 @@ function App() {
         errors: [],
         defaults: true,
       });
-      setInputData({ data: null, errors: [] });
+      setInputData({ data: defaults(schemas.input), errors: [] });
     }
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [schemas.config]);
@@ -286,7 +400,7 @@ function App() {
               <div className="flex flex-col gap-3">
                 <h2 className="text-xl font-semibold">Output</h2>
                 <div className="p-4 border border-divider-700 flex flex-col gap-3 rounded-2xl bg-background text-lg">
-                  {latest.streamed_output.map(str).join("") || "..."}
+                  <StreamOutput streamed={latest.streamed_output} />
                 </div>
                 <IntermediateSteps latest={latest} />
               </div>

langserve/playground/src/components/ChatMessageTuplesControlRenderer.tsx

@@ -0,0 +1,111 @@
+import { withJsonFormsControlProps } from "@jsonforms/react";
+import PlusIcon from "../assets/PlusIcon.svg?react";
+import TrashIcon from "../assets/TrashIcon.svg?react";
+import {
+  rankWith,
+  and,
+  schemaMatches,
+  Paths,
+  isControl,
+} from "@jsonforms/core";
+import { AutosizeTextarea } from "./AutosizeTextarea";
+
+type MessageTuple = [string, string];
+
+export const chatMessagesTupleTester = rankWith(
+  12,
+  and(
+    isControl,
+    schemaMatches((schema) => {
+      if (schema.type !== "array") return false;
+      if (typeof schema.items !== "object" || schema.items == null)
+        return false;
+
+      if ("type" in schema.items) {
+        return (
+          schema.items.type === "array" &&
+          schema.items.minItems === 2 &&
+          schema.items.maxItems === 2 &&
+          Array.isArray(schema.items.items) &&
+          schema.items.items.length === 2 &&
+          schema.items.items.every((schema) => schema.type === "string") &&
+          schema.title === "Chat History"
+        );
+      }
+
+      return false;
+    })
+  )
+);
+
+export const ChatMessageTuplesControlRenderer = withJsonFormsControlProps(
+  (props) => {
+    const data: Array<MessageTuple> = props.data ?? [];
+
+    return (
+      <div className="control">
+        <div className="flex items-center justify-between">
+          <label className="text-xs uppercase font-semibold text-ls-gray-100">
+            {props.label}
+          </label>
+          <button
+            className="p-1 rounded-full"
+            onClick={() => {
+              const lastRole = data.length ? data[data.length - 1][0] : "ai";
+              props.handleChange(props.path, [
+                ...data,
+                [lastRole === "ai" ? "human" : "ai", ""],
+              ]);
+            }}
+          >
+            <PlusIcon className="w-5 h-5" />
+          </button>
+        </div>
+
+        <div className="flex flex-col gap-3 mt-1 empty:hidden">
+          {data.map(([type, content], index) => {
+            const msgPath = Paths.compose(props.path, `${index}`);
+            return (
+              <div className="control group" key={index}>
+                <div className="flex items-start justify-between gap-2">
+                  <select
+                    className="-ml-1 min-w-[100px]"
+                    value={type}
+                    onChange={(e) => {
+                      props.handleChange(
+                        Paths.compose(msgPath, "0"),
+                        e.target.value
+                      );
+                    }}
+                  >
+                    <option value="human">Human</option>
+                    <option value="ai">AI</option>
+                    <option value="system">System</option>
+                  </select>
+                  <button
+                    className="p-1 border rounded opacity-0 transition-opacity border-divider-700 group-focus-within:opacity-100 group-hover:opacity-100"
+                    onClick={() => {
+                      props.handleChange(
+                        props.path,
+                        data.filter((_, i) => i !== index)
+                      );
+                    }}
+                  >
+                    <TrashIcon className="w-4 h-4" />
+                  </button>
+                </div>
+
+                <AutosizeTextarea
+                  value={content}
+                  onChange={(content) => {
+                    props.handleChange(Paths.compose(msgPath, "1"), content);
+                  }}
+                />
+              </div>
+            );
+          })}
+        </div>
+      </div>
+    );
+  }
+);

langserve/playground/src/components/ChatMessagesControlRenderer.tsx

@@ -0,0 +1,163 @@
+import { withJsonFormsControlProps } from "@jsonforms/react";
+import PlusIcon from "../assets/PlusIcon.svg?react";
+import TrashIcon from "../assets/TrashIcon.svg?react";
+import {
+  rankWith,
+  and,
+  schemaMatches,
+  Paths,
+  isControl,
+} from "@jsonforms/core";
+import { AutosizeTextarea } from "./AutosizeTextarea";
+import _ from "lodash";
+
+export const chatMessagesTester = rankWith(
+  12,
+  and(
+    isControl,
+    schemaMatches((schema) => {
+      if (schema.type !== "array") return false;
+      if (typeof schema.items !== "object" || schema.items == null)
+        return false;
+
+      if (
+        "type" in schema.items &&
+        schema.items.type != null &&
+        schema.items.title != null
+      ) {
+        return (
+          schema.items.type === "object" &&
+          (schema.items.title?.endsWith("Message") ||
+            schema.items.title?.endsWith("MessageChunk"))
+        );
+      }
+
+      if ("anyOf" in schema.items && schema.items.anyOf != null) {
+        return schema.items.anyOf.every(
+          (schema) =>
+            schema.type === "object" &&
+            (schema.title?.endsWith("Message") ||
+              schema.title?.endsWith("MessageChunk"))
+        );
+      }
+
+      return false;
+    })
+  )
+);
+
+interface MessageFields {
+  content: string;
+  additional_kwargs?: { [key: string]: unknown };
+  name?: string;
+  type?: string;
+
+  role?: string;
+}
+
+export const ChatMessagesControlRenderer = withJsonFormsControlProps(
+  (props) => {
+    const data: Array<MessageFields> = props.data ?? [];
+
+    return (
+      <div className="control">
+        <div className="flex items-center justify-between">
+          <label className="text-xs uppercase font-semibold text-ls-gray-100">
+            {props.label}
+          </label>
+          <button
+            className="p-1 rounded-full"
+            onClick={() => {
+              const lastRole = data.length ? data[data.length - 1].type : "ai";
+              props.handleChange(props.path, [
+                ...data,
+                { content: "", type: lastRole === "human" ? "ai" : "human" },
+              ]);
+            }}
+          >
+            <PlusIcon className="w-5 h-5" />
+          </button>
+        </div>
+
+        <div className="flex flex-col gap-3 mt-1 empty:hidden">
+          {data.map((message, index) => {
+            const msgPath = Paths.compose(props.path, `${index}`);
+            const type = message.type ?? "chat";
+            return (
+              <div className="control group" key={index}>
+                <div className="flex items-start justify-between gap-2">
+                  <select
+                    className="-ml-1 min-w-[100px]"
+                    value={type}
+                    onChange={(e) => {
+                      props.handleChange(
+                        Paths.compose(msgPath, "type"),
+                        e.target.value
+                      );
+                    }}
+                  >
+                    <option value="human">Human</option>
+                    <option value="ai">AI</option>
+                    <option value="system">System</option>
+                    <option value="function">Function</option>
+
+                    <option value="chat">Chat</option>
+                  </select>
+                  <button
+                    className="p-1 border rounded opacity-0 transition-opacity border-divider-700 group-focus-within:opacity-100 group-hover:opacity-100"
+                    onClick={() => {
+                      props.handleChange(
+                        props.path,
+                        data.filter((_, i) => i !== index)
+                      );
+                    }}
+                  >
+                    <TrashIcon className="w-4 h-4" />
+                  </button>
+                </div>
+
+                {type === "chat" && (
+                  <input
+                    className="mb-1"
+                    placeholder="Role"
+                    value={message.role ?? ""}
+                    onChange={(e) => {
+                      props.handleChange(
+                        Paths.compose(msgPath, "role"),
+                        e.target.value
+                      );
+                    }}
+                  />
+                )}
+
+                {type === "function" && (
+                  <input
+                    className="mb-1"
+                    placeholder="Function Name"
+                    value={message.name ?? ""}
+                    onChange={(e) => {
+                      props.handleChange(
+                        Paths.compose(msgPath, "name"),
+                        e.target.value
+                      );
+                    }}
+                  />
+                )}
+
+                <AutosizeTextarea
+                  value={message.content}
+                  onChange={(content) => {
+                    props.handleChange(
+                      Paths.compose(msgPath, "content"),
+                      content
+                    );
+                  }}
+                />
+              </div>
+            );
+          })}
+        </div>
+      </div>
+    );
+  }
+);

langserve/playground/src/components/CustomArrayControlRenderer.tsx

@@ -84,7 +84,7 @@ export const MaterialArrayControlRenderer = (props: ArrayLayoutProps) => {
 
 // eslint-disable-next-line react-refresh/only-export-components
 export const materialArrayControlTester: RankedTester = rankWith(
-  999,
+  11,
   or(isObjectArrayControl, isPrimitiveArrayControl, isObjectArrayWithNesting)
 );
 

langserve/playground/src/useSchemas.tsx

@@ -55,7 +55,7 @@ export function useSchemas(
     if (!debouncedConfigData.defaults) {
       fetch(
         resolveApiUrl(
-          `c/${compressToEncodedURIComponent(
+          `/c/${compressToEncodedURIComponent(
             JSON.stringify(debouncedConfigData.data)
           )}/input_schema`
         )

langserve/serialization.py

@@ -1,10 +1,20 @@
-"""Serialization module for Well Known LangChain objects.
+"""Serialization for well known objects and callback events.
 
 Specialized JSON serialization for well known LangChain objects that
 can be expected to be frequently transmitted between chains.
+
+Callback events handle well known objects together with a few other
+common types like UUIDs and Exceptions that might appear in the callback.
+
+By default, exceptions are serialized as a generic exception without
+any information about the exception. This is done to prevent leaking
+sensitive information from the server to the client.
 """
+import abc
 import json
-from typing import Any, Union
+import logging
+from functools import lru_cache
+from typing import Any, Dict, List, Union
 
 from langchain.prompts.base import StringPromptValue
 from langchain.prompts.chat import ChatPromptValueConcrete
@@ -22,6 +32,14 @@ from langchain.schema.messages import (
     SystemMessage,
     SystemMessageChunk,
 )
+from langchain.schema.output import (
+    ChatGeneration,
+    ChatGenerationChunk,
+    Generation,
+    LLMResult,
+)
+
+from langserve.validation import CallbackEvent
 
 try:
     from pydantic.v1 import BaseModel, ValidationError
@@ -29,6 +47,15 @@ except ImportError:
     from pydantic import BaseModel, ValidationError
 
 
+logger = logging.getLogger(__name__)
+
+
+@lru_cache(maxsize=1_000)  # Will accommodate up to 1_000 different error messages
+def _log_error_message_once(error_message: str) -> None:
+    """Log an error message once."""
+    logger.error(error_message)
+
+
 class WellKnownLCObject(BaseModel):
     """A well known LangChain object.
 
@@ -54,6 +81,10 @@ class WellKnownLCObject(BaseModel):
         AgentAction,
         AgentFinish,
         AgentActionMessageLog,
+        LLMResult,
+        ChatGeneration,
+        Generation,
+        ChatGenerationChunk,
     ]
 
 
@@ -67,41 +98,141 @@ class _LangChainEncoder(json.JSONEncoder):
         return super().default(obj)
 
 
-# Custom JSON Decoder
-class _LangChainDecoder(json.JSONDecoder):
-    """Custom JSON Decoder that handles well known LangChain objects."""
+def _decode_lc_objects(value: Any) -> Any:
+    """Decode the value."""
+    if isinstance(value, dict):
+        v = {key: _decode_lc_objects(v) for key, v in value.items()}
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        """Initialize the LangChainDecoder."""
-        super().__init__(object_hook=self.decoder, *args, **kwargs)
+        try:
+            obj = WellKnownLCObject.parse_obj(v)
+            parsed = obj.__root__
+            if set(parsed.dict()) != set(value):
+                raise ValueError("Invalid object")
+            return parsed
+        except (ValidationError, ValueError):
+            return v
+    elif isinstance(value, list):
+        return [_decode_lc_objects(item) for item in value]
+    else:
+        return value
 
-    def decoder(self, value) -> Any:
-        """Decode the value."""
-        if isinstance(value, dict):
+
+class ServerSideException(Exception):
+    """Exception raised when a server side exception occurs.
+
+    The goal of this exception is to provide a way to communicate
+    to the client that a server side exception occurred without
+    revealing too much information about the exception as it may contain
+    sensitive information.
+    """
+
+
+def _decode_event_data(value: Any) -> Any:
+    """Decode the event data from a JSON object representation."""
+    if isinstance(value, dict):
+        try:
+            obj = CallbackEvent.parse_obj(value)
+            return obj.__root__
+        except ValidationError:
             try:
                 obj = WellKnownLCObject.parse_obj(value)
                 return obj.__root__
             except ValidationError:
-                return {key: self.decoder(v) for key, v in value.items()}
-        elif isinstance(value, list):
-            return [self.decoder(item) for item in value]
-        else:
-            return value
+                return {key: _decode_event_data(v) for key, v in value.items()}
+    elif isinstance(value, list):
+        return [_decode_event_data(item) for item in value]
+    else:
+        return value
 
 
 # PUBLIC API
 
 
-def simple_dumpd(obj: Any) -> Any:
-    """Convert the given object to a JSON serializable object."""
-    return json.loads(json.dumps(obj, cls=_LangChainEncoder))
+class Serializer(abc.ABC):
+    @abc.abstractmethod
+    def dumpd(self, obj: Any) -> Any:
+        """Convert the given object to a JSON serializable object."""
+
+    @abc.abstractmethod
+    def dumps(self, obj: Any) -> str:
+        """Dump the given object as a JSON string."""
+
+    @abc.abstractmethod
+    def loads(self, s: str) -> Any:
+        """Load the given JSON string."""
+
+    @abc.abstractmethod
+    def loadd(self, obj: Any) -> Any:
+        """Load the given object."""
 
 
-def simple_dumps(obj: Any) -> str:
-    """Dump the given object as a JSON string."""
-    return json.dumps(obj, cls=_LangChainEncoder)
+class WellKnownLCSerializer(Serializer):
+    def dumpd(self, obj: Any) -> Any:
+        """Convert the given object to a JSON serializable object."""
+        return json.loads(json.dumps(obj, cls=_LangChainEncoder))  # :*(
+
+    def dumps(self, obj: Any) -> str:
+        """Dump the given object as a JSON string."""
+        return json.dumps(obj, cls=_LangChainEncoder)
+
+    def loadd(self, obj: Any) -> Any:
+        return _decode_lc_objects(obj)
+
+    def loads(self, s: str) -> Any:
+        """Load the given JSON string."""
+        return self.loadd(json.loads(s))
 
 
-def simple_loads(s: str) -> Any:
-    """Load the given JSON string."""
-    return json.loads(s, cls=_LangChainDecoder)
+def _project_top_level(model: BaseModel) -> Dict[str, Any]:
+    """Project the top level of the model as dict."""
+    return {key: getattr(model, key) for key in model.__fields__}
+
+
+def load_events(events: Any) -> List[Dict[str, Any]]:
+    """Load and validate the event.
+
+    Args:
+        events: The events to load and validate.
+
+    Returns:
+        The loaded and validated events.
+    """
+    if not isinstance(events, list):
+        _log_error_message_once(f"Expected a list got {type(events)}")
+        return []
+
+    decoded_events = []
+
+    for event in events:
+        if not isinstance(event, dict):
+            _log_error_message_once(f"Expected a dict got {type(event)}")
+            # Discard the event / potentially error
+            continue
+
+        # First load all inner objects
+        decoded_event_data = {
+            key: _decode_lc_objects(value) for key, value in event.items()
+        }
+
+        # Then validate the event
+        try:
+            full_event = CallbackEvent.parse_obj(decoded_event_data)
+        except ValidationError as e:
+            msg = f"Encountered an invalid event: {e}"
+            if "type" in decoded_event_data:
+                msg += f' of type {repr(decoded_event_data["type"])}'
+            _log_error_message_once(msg)
+            continue
+
+        decoded_event_data = _project_top_level(full_event.__root__)
+
+        if decoded_event_data["type"].endswith("_error"):
+            # Data is validated by this point, so we can assume that the shape
+            # of the data is correct
+            error = decoded_event_data["error"]
+            msg = f"{error['status_code']}: {error['message']}"
+            decoded_event_data["error"] = ServerSideException(msg)
+
+        decoded_events.append(decoded_event_data)
+
+    return decoded_events

langserve/server.py

@@ -24,11 +24,11 @@ from fastapi import HTTPException, Request
 from langchain.callbacks.tracers.log_stream import RunLogPatch
 from langchain.load.serializable import Serializable
 from langchain.schema.runnable import Runnable
-from langchain.schema.runnable.config import merge_configs
+from langchain.schema.runnable.config import get_config_list, merge_configs
 from typing_extensions import Annotated
 
+from langserve.callbacks import AsyncEventAggregatorCallback, CallbackEventDict
 from langserve.lzstring import LZString
-from langserve.version import __version__
 
 try:
     from pydantic.v1 import BaseModel, create_model
@@ -36,7 +36,7 @@ except ImportError:
     from pydantic import BaseModel, Field, create_model
 
 from langserve.playground import serve_playground
-from langserve.serialization import simple_dumpd, simple_dumps
+from langserve.serialization import WellKnownLCSerializer
 from langserve.validation import (
     create_batch_request_model,
     create_batch_response_model,
@@ -45,6 +45,7 @@ from langserve.validation import (
     create_stream_log_request_model,
     create_stream_request_model,
 )
+from langserve.version import __version__
 
 try:
     from fastapi import APIRouter, FastAPI
@@ -199,7 +200,42 @@ def _add_tracing_info_to_metadata(config: Dict[str, Any], request: Request) -> N
     config["metadata"] = metadata
 
 
+def _scrub_exceptions_in_event(event: CallbackEventDict) -> CallbackEventDict:
+    """Scrub exceptions and change to a serializable format."""
+    type_ = event["type"]
+    # Check if the event type is one that could contain an error key
+    # for example, on_chain_error, on_tool_error, etc.
+    if "error" not in type_:
+        return event
+
+    # This is not scrubbing -- it's doing serialization
+    if "error" not in event:  # if there is an error key, scrub it
+        return event
+
+    if isinstance(event["error"], BaseException):
+        event.copy()
+        event["error"] = {"status_code": 500, "message": "Internal Server Error"}
+        return event
+
+    raise AssertionError(f"Expected an exception got {type(event['error'])}")
+
+
 _APP_SEEN = weakref.WeakSet()
+_APP_TO_PATHS = weakref.WeakKeyDictionary()
+
+
+def _register_path_for_app(app: Union[FastAPI, APIRouter], path: str) -> None:
+    """Register a path when its added to app. Raise if path already seen."""
+    if app in _APP_TO_PATHS:
+        seen_paths = _APP_TO_PATHS.get(app)
+        if path in seen_paths:
+            raise ValueError(
+                f"A runnable already exists at path: {path}. If adding "
+                f"multiple runnables make sure they have different paths."
+            )
+        seen_paths.add(path)
+    else:
+        _APP_TO_PATHS[app] = {path}
 
 
 # PUBLIC API
@@ -213,6 +249,7 @@ def add_routes(
     input_type: Union[Type, Literal["auto"], BaseModel] = "auto",
     output_type: Union[Type, Literal["auto"], BaseModel] = "auto",
     config_keys: Sequence[str] = (),
+    include_callback_events: bool = False,
 ) -> None:
     """Register the routes on the given FastAPI app or APIRouter.
 
@@ -234,11 +271,19 @@ def add_routes(
         input_type: type to use for input validation.
             Default is "auto" which will use the InputType of the runnable.
             User is free to provide a custom type annotation.
+            Favor using runnable.with_types(input_type=..., output_type=...) instead.
+            This parameter may get deprecated!
         output_type: type to use for output validation.
             Default is "auto" which will use the OutputType of the runnable.
             User is free to provide a custom type annotation.
+            Favor using runnable.with_types(input_type=..., output_type=...) instead.
+            This parameter may get deprecated!
         config_keys: list of config keys that will be accepted, by default
                      no config keys are accepted.
+        include_callback_events: Whether to include callback events in the response.
+            If true, the client will be able to show trace information
+            including events that occurred on the server side.
+            Be sure not to include any sensitive information in the callback events.
     """
     try:
         from sse_starlette import EventSourceResponse
@@ -249,6 +294,9 @@ def add_routes(
             "Use `pip install sse_starlette` to install."
         )
 
+    _register_path_for_app(app, path)
+    well_known_lc_serializer = WellKnownLCSerializer()
+
     if hasattr(app, "openapi_tags") and app not in _APP_SEEN:
         _APP_SEEN.add(app)
         app.openapi_tags = [
@@ -258,22 +306,37 @@ def add_routes(
             },
             {
                 "name": "config",
-                "description": "Endpoints with a default configuration set by `config_hash` path parameter.",  # noqa: E501
+                "description": (
+                    "Endpoints with a default configuration "
+                    "set by `config_hash` path parameter."
+                ),
             },
         ]
 
+    if path and not path.startswith("/"):
+        raise ValueError(
+            f"Got an invalid path: {path}. "
+            f"If specifying path please start it with a `/`"
+        )
+
     namespace = path or ""
 
     model_namespace = _replace_non_alphanumeric_with_underscores(path.strip("/"))
 
-    input_type_ = _resolve_model(
-        runnable.input_schema if input_type == "auto" else input_type,
-        "Input",
-        model_namespace,
-    )
+    with_types = {}
+
+    if input_type != "auto":
+        with_types["input_type"] = input_type
+    if output_type != "auto":
+        with_types["output_type"] = output_type
+
+    if with_types:
+        runnable = runnable.with_types(**with_types)
+
+    input_type_ = _resolve_model(runnable.get_input_schema(), "Input", model_namespace)
 
     output_type_ = _resolve_model(
-        runnable.output_schema if output_type == "auto" else output_type,
+        runnable.get_output_schema(),
         "Output",
         model_namespace,
     )
@@ -281,6 +344,7 @@ def add_routes(
     ConfigPayload = _add_namespace_to_model(
         model_namespace, runnable.config_schema(include=config_keys)
     )
+
     InvokeRequest = create_invoke_request_model(
         model_namespace, input_type_, ConfigPayload
     )
@@ -315,11 +379,27 @@ def add_routes(
             config_hash, invoke_request.config, keys=config_keys, model=ConfigPayload
         )
         _add_tracing_info_to_metadata(config, request)
+        event_aggregator = AsyncEventAggregatorCallback()
+        config["callbacks"] = [event_aggregator]
         output = await runnable.ainvoke(
-            _unpack_input(invoke_request.input), config=config
+            _unpack_input(invoke_request.input),
+            config=config,
         )
 
-        return InvokeResponse(output=simple_dumpd(output))
+        if include_callback_events:
+            callback_events = [
+                _scrub_exceptions_in_event(event)
+                for event in event_aggregator.callback_events
+            ]
+        else:
+            callback_events = []
+
+        return InvokeResponse(
+            output=well_known_lc_serializer.dumpd(output),
+            # Callbacks are scrubbed and exceptions are converted to serializable format
+            # before returned in the response.
+            callback_events=callback_events,
+        )
 
     @app.post(
         namespace + "/c/{config_hash}/batch",
@@ -333,25 +413,57 @@ def add_routes(
         config_hash: str = "",
     ) -> BatchResponse:
         """Invoke the runnable with the given inputs and config."""
+        # First convert to list type
         if isinstance(batch_request.config, list):
-            config = [
+            configs = [
                 _unpack_config(
                     config_hash, config, keys=config_keys, model=ConfigPayload
                 )
                 for config in batch_request.config
             ]
-
-            for c in config:
-                _add_tracing_info_to_metadata(c, request)
         else:
-            config = _unpack_config(
-                config_hash, batch_request.config, keys=config_keys, model=ConfigPayload
+            configs = _unpack_config(
+                config_hash,
+                batch_request.config,
+                keys=config_keys,
+                model=ConfigPayload,
             )
-            _add_tracing_info_to_metadata(config, request)
-        inputs = [_unpack_input(input_) for input_ in batch_request.inputs]
-        output = await runnable.abatch(inputs, config=config)
 
-        return BatchResponse(output=simple_dumpd(output))
+        # Unpack
+
+        # Make sure that the number of configs matches the number of inputs
+        # Since we'll be adding callbacks to the configs.
+        _configs = get_config_list(configs, len(batch_request.inputs))
+
+        aggregators = [
+            AsyncEventAggregatorCallback() for _ in range(len(batch_request.inputs))
+        ]
+
+        for c, aggregator in zip(_configs, aggregators):
+            _add_tracing_info_to_metadata(c, request)
+            c["callbacks"] = [aggregator]
+
+        inputs = [_unpack_input(input_) for input_ in batch_request.inputs]
+
+        output = await runnable.abatch(inputs, config=_configs)
+
+        if include_callback_events:
+            callback_events = [
+                # Scrub sensitive information and convert
+                # exceptions to serializable format
+                [
+                    _scrub_exceptions_in_event(event)
+                    for event in aggregator.callback_events
+                ]
+                for aggregator in aggregators
+            ]
+        else:
+            callback_events = []
+
+        return BatchResponse(
+            output=well_known_lc_serializer.dumpd(output),
+            callback_events=callback_events,
+        )
 
     @app.post(namespace + "/c/{config_hash}/stream", tags=["config"])
     @app.post(f"{namespace}/stream")
@@ -417,7 +529,10 @@ def add_routes(
                     input_,
                     config=config,
                 ):
-                    yield {"data": simple_dumps(chunk), "event": "data"}
+                    yield {
+                        "data": well_known_lc_serializer.dumps(chunk),
+                        "event": "data",
+                    }
                 yield {"event": "end"}
             except BaseException:
                 yield {
@@ -518,7 +633,7 @@ def add_routes(
 
                     # Temporary adapter
                     yield {
-                        "data": simple_dumps(data),
+                        "data": well_known_lc_serializer.dumps(data),
                         "event": "data",
                     }
                 yield {"event": "end"}
@@ -540,37 +655,24 @@ def add_routes(
     @app.get(f"{namespace}/input_schema")
     async def input_schema(config_hash: str = "") -> Any:
         """Return the input schema of the runnable."""
-        return (
-            runnable.with_config(
-                _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
-            ).input_schema.schema()
-            if input_type == "auto"
-            else input_type_.schema()
-        )
+        return runnable.get_input_schema(
+            _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
+        ).schema()
 
     @app.get(namespace + "/c/{config_hash}/output_schema", tags=["config"])
     @app.get(f"{namespace}/output_schema")
     async def output_schema(config_hash: str = "") -> Any:
         """Return the output schema of the runnable."""
-        return (
-            runnable.with_config(
-                _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
-            ).output_schema.schema()
-            if output_type_ == "auto"
-            else output_type_.schema()
-        )
+        return runnable.get_output_schema(
+            _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
+        ).schema()
 
     @app.get(namespace + "/c/{config_hash}/config_schema", tags=["config"])
     @app.get(f"{namespace}/config_schema")
     async def config_schema(config_hash: str = "") -> Any:
         """Return the config schema of the runnable."""
-        return (
-            runnable.with_config(
-                _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
-            )
-            .config_schema(include=config_keys)
-            .schema()
-        )
+        config = _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
+        return runnable.with_config(config).config_schema(include=config_keys).schema()
 
     @app.get(
         namespace + "/c/{config_hash}/playground/{file_path:path}",
@@ -580,15 +682,10 @@ def add_routes(
     @app.get(namespace + "/playground/{file_path:path}", include_in_schema=False)
     async def playground(file_path: str, config_hash: str = "") -> Any:
         """Return the playground of the runnable."""
+        config = _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
         return await serve_playground(
-            runnable.with_config(
-                _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
-            ),
-            runnable.with_config(
-                _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
-            ).input_schema
-            if input_type == "auto"
-            else input_type_,
+            runnable.with_config(config),
+            runnable.with_config(config).input_schema,
             config_keys,
             f"{namespace}/playground",
             file_path,

langserve/validation.py

@@ -16,7 +16,16 @@ Models are created with a namespace to avoid name collisions when hosting
 multiple runnables. When present the name collisions prevent fastapi from
 generating OpenAPI specs.
 """
-from typing import List, Optional, Sequence, Union
+from typing import Any, Dict, List, Literal, Optional, Sequence, Union
+from uuid import UUID
+
+from langchain.schema import (
+    BaseMessage,
+    ChatGeneration,
+    Document,
+    Generation,
+    RunInfo,
+)
 
 try:
     from pydantic.v1 import BaseModel, Field, create_model
@@ -194,6 +203,10 @@ def create_invoke_response_model(
     invoke_response_type = create_model(
         f"{namespace}InvokeResponse",
         output=(output_type, Field(..., description="The output of the invocation.")),
+        callback_events=(
+            List[CallbackEvent],
+            Field(..., description="Callback events generated by the server side."),
+        ),
     )
     invoke_response_type.update_forward_refs()
     return invoke_response_type
@@ -217,6 +230,207 @@ def create_batch_response_model(
                 ),
             ),
         ),
+        callback_events=(
+            List[List[CallbackEvent]],
+            Field(
+                ...,
+                description=(
+                    "Callback events generated by the server side."
+                    "The outer list corresponds to the inputs and the inner "
+                    "list corresponds to the callbacks generated for that input."
+                ),
+            ),
+        ),
     )
     batch_response_type.update_forward_refs()
     return batch_response_type
+
+
+# Pydantic validators for callback events
+# These objects may have a slightly different shape than the callback events
+# used internally in langchain because they represent a serialized version
+# of the callback event.
+# For example, exceptions are replaced by error objects consisting of a
+# status code and a message.
+
+
+class OnChainStart(BaseModel):
+    """On Chain Start Callback Event."""
+
+    serialized: Dict[str, Any]
+    inputs: Any
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    kwargs: Any = None
+    type: Literal["on_chain_start"] = "on_chain_start"
+
+
+class OnChainEnd(BaseModel):
+    """On Chain End Callback Event."""
+
+    outputs: Any
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    kwargs: Any = None
+    type: Literal["on_chain_end"] = "on_chain_end"
+
+
+class Error(BaseModel):
+    """Error object that is modeled after an HTTP error format."""
+
+    status_code: int
+    message: str
+    type: Literal["error"] = "error"
+
+
+class OnChainError(BaseModel):
+    """On Chain Error Callback Event."""
+
+    error: Error
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    kwargs: Any = None
+    type: Literal["on_chain_error"] = "on_chain_error"
+
+
+class OnToolStart(BaseModel):
+    """On Tool Start Callback Event."""
+
+    serialized: Dict[str, Any]
+    input_str: str
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    kwargs: Any = None
+    type: Literal["on_tool_start"] = "on_tool_start"
+
+
+class OnToolEnd(BaseModel):
+    """On Tool End Callback Event."""
+
+    output: str
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    kwargs: Any = None
+    type: Literal["on_tool_end"] = "on_tool_end"
+
+
+class OnToolError(BaseModel):
+    """On Tool Error Callback Event."""
+
+    error: Error
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    kwargs: Any = None
+    type: Literal["on_tool_error"] = "on_tool_error"
+
+
+class OnChatModelStart(BaseModel):
+    """On Chat Model Start Callback Event."""
+
+    serialized: Dict[str, Any]
+    messages: List[List[BaseMessage]]
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    kwargs: Any = None
+    type: Literal["on_chat_model_start"] = "on_chat_model_start"
+
+
+class OnLLMStart(BaseModel):
+    """On LLM Start Callback Event."""
+
+    serialized: Dict[str, Any]
+    prompts: List[str]
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    kwargs: Any = None
+    type: Literal["on_llm_start"] = "on_llm_start"
+
+
+class LLMResult(BaseModel):
+    """Concrete instance of LLMResult for validation only.
+
+    Must be kept in sync with langchain.schema.llm.LLMResult.
+    """
+
+    generations: List[List[Union[Generation, ChatGeneration]]]
+    """List of generated outputs. This is a List[List[]] because
+    each input could have multiple candidate generations."""
+    llm_output: Optional[dict] = None
+    """Arbitrary LLM provider-specific output."""
+    run: Optional[List[RunInfo]] = None
+    """List of metadata info for model call for each input."""
+
+
+class OnLLMEnd(BaseModel):
+    """On LLM End Callback Event."""
+
+    response: LLMResult
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    kwargs: Any = None
+    type: Literal["on_llm_end"] = "on_llm_end"
+
+
+class OnRetrieverStart(BaseModel):
+    """On Retriever Start Callback Event."""
+
+    serialized: Dict[str, Any]
+    query: str
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    kwargs: Any = None
+    type: Literal["on_retriever_start"] = "on_retriever_start"
+
+
+class OnRetrieverError(BaseModel):
+    """On Retriever Error Callback Event."""
+
+    error: Error
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    kwargs: Any = None
+    type: Literal["on_retriever_error"] = "on_retriever_error"
+
+
+class OnRetrieverEnd(BaseModel):
+    """On Retriever End Callback Event."""
+
+    documents: Sequence[Document]
+    run_id: UUID
+    parent_run_id: Optional[UUID] = None
+    tags: Optional[List[str]] = None
+    kwargs: Any = None
+    type: Literal["on_retriever_end"] = "on_retriever_end"
+
+
+class CallbackEvent(BaseModel):
+    __root__: Union[
+        OnChainStart,
+        OnChainEnd,
+        OnChainError,
+        OnChatModelStart,
+        OnLLMStart,
+        OnLLMEnd,
+        OnToolStart,
+        OnToolEnd,
+        OnToolError,
+        OnRetrieverStart,
+        OnRetrieverEnd,
+        OnRetrieverError,
+    ]

poetry.lock

@@ -966,7 +966,7 @@ files = [
     {file = "greenlet-3.0.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:0b72b802496cccbd9b31acea72b6f87e7771ccfd7f7927437d592e5c92ed703c"},
     {file = "greenlet-3.0.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:527cd90ba3d8d7ae7dceb06fda619895768a46a1b4e423bdb24c1969823b8362"},
     {file = "greenlet-3.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:37f60b3a42d8b5499be910d1267b24355c495064f271cfe74bf28b17b099133c"},
-    {file = "greenlet-3.0.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:1482fba7fbed96ea7842b5a7fc11d61727e8be75a077e603e8ab49d24e234383"},
+    {file = "greenlet-3.0.0-cp311-universal2-macosx_10_9_universal2.whl", hash = "sha256:c3692ecf3fe754c8c0f2c95ff19626584459eab110eaab66413b1e7425cd84e9"},
     {file = "greenlet-3.0.0-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:be557119bf467d37a8099d91fbf11b2de5eb1fd5fc5b91598407574848dc910f"},
     {file = "greenlet-3.0.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:73b2f1922a39d5d59cc0e597987300df3396b148a9bd10b76a058a2f2772fc04"},
     {file = "greenlet-3.0.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d1e22c22f7826096ad503e9bb681b05b8c1f5a8138469b255eb91f26a76634f2"},
@@ -976,6 +976,7 @@ files = [
     {file = "greenlet-3.0.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:952256c2bc5b4ee8df8dfc54fc4de330970bf5d79253c863fb5e6761f00dda35"},
     {file = "greenlet-3.0.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:269d06fa0f9624455ce08ae0179430eea61085e3cf6457f05982b37fd2cefe17"},
     {file = "greenlet-3.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:9adbd8ecf097e34ada8efde9b6fec4dd2a903b1e98037adf72d12993a1c80b51"},
+    {file = "greenlet-3.0.0-cp312-universal2-macosx_10_9_universal2.whl", hash = "sha256:553d6fb2324e7f4f0899e5ad2c427a4579ed4873f42124beba763f16032959af"},
     {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c6b5ce7f40f0e2f8b88c28e6691ca6806814157ff05e794cdd161be928550f4c"},
     {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ecf94aa539e97a8411b5ea52fc6ccd8371be9550c4041011a091eb8b3ca1d810"},
     {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:80dcd3c938cbcac986c5c92779db8e8ce51a89a849c135172c88ecbdc8c056b7"},
@@ -1654,13 +1655,13 @@ test = ["hatch", "ipykernel", "openapi-core (>=0.18.0,<0.19.0)", "openapi-spec-v
 
 [[package]]
 name = "langchain"
-version = "0.0.319"
+version = "0.0.322"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = ">=3.8.1,<4.0"
 files = [
-    {file = "langchain-0.0.319-py3-none-any.whl", hash = "sha256:a61448fd418ff9478f2be3477c9c92acbf6b6acd8163c58c994a6158d4d116f8"},
-    {file = "langchain-0.0.319.tar.gz", hash = "sha256:4fe5025e5fd48dcf8e02107fefe173ba997af3c8960871cc4a4467e24bb89375"},
+    {file = "langchain-0.0.322-py3-none-any.whl", hash = "sha256:f065d19eff2702cd941fd1f6adca6baf2a1181387f04d3429a0ec7a197e96155"},
+    {file = "langchain-0.0.322.tar.gz", hash = "sha256:8415327a964bff9d643202697aca1dcbcfd2d89fb0b49f799bdea37d01fa7f51"},
 ]
 
 [package.dependencies]
@@ -1734,6 +1735,16 @@ files = [
     {file = "MarkupSafe-2.1.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:5bbe06f8eeafd38e5d0a4894ffec89378b6c6a625ff57e3028921f8ff59318ac"},
     {file = "MarkupSafe-2.1.3-cp311-cp311-win32.whl", hash = "sha256:dd15ff04ffd7e05ffcb7fe79f1b98041b8ea30ae9234aed2a9168b5797c3effb"},
     {file = "MarkupSafe-2.1.3-cp311-cp311-win_amd64.whl", hash = "sha256:134da1eca9ec0ae528110ccc9e48041e0828d79f24121a1a146161103c76e686"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:f698de3fd0c4e6972b92290a45bd9b1536bffe8c6759c62471efaa8acb4c37bc"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:aa57bd9cf8ae831a362185ee444e15a93ecb2e344c8e52e4d721ea3ab6ef1823"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ffcc3f7c66b5f5b7931a5aa68fc9cecc51e685ef90282f4a82f0f5e9b704ad11"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47d4f1c5f80fc62fdd7777d0d40a2e9dda0a05883ab11374334f6c4de38adffd"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1f67c7038d560d92149c060157d623c542173016c4babc0c1913cca0564b9939"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:9aad3c1755095ce347e26488214ef77e0485a3c34a50c5a5e2471dff60b9dd9c"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:14ff806850827afd6b07a5f32bd917fb7f45b046ba40c57abdb636674a8b559c"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8f9293864fe09b8149f0cc42ce56e3f0e54de883a9de90cd427f191c346eb2e1"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-win32.whl", hash = "sha256:715d3562f79d540f251b99ebd6d8baa547118974341db04f5ad06d5ea3eb8007"},
+    {file = "MarkupSafe-2.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:1b8dd8c3fd14349433c79fa8abeb573a55fc0fdd769133baac1f5e07abf54aeb"},
     {file = "MarkupSafe-2.1.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2"},
     {file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cb0932dc158471523c9637e807d9bfb93e06a95cbf010f1a38b98623b929ef2b"},
     {file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707"},
@@ -2438,13 +2449,13 @@ plugins = ["importlib-metadata"]
 
 [[package]]
 name = "pytest"
-version = "7.4.2"
+version = "7.4.3"
 description = "pytest: simple powerful testing with Python"
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "pytest-7.4.2-py3-none-any.whl", hash = "sha256:1d881c6124e08ff0a1bb75ba3ec0bfd8b5354a01c194ddd5a0a870a48d99b002"},
-    {file = "pytest-7.4.2.tar.gz", hash = "sha256:a766259cfab564a2ad52cb1aae1b881a75c3eb7e34ca3779697c23ed47c47069"},
+    {file = "pytest-7.4.3-py3-none-any.whl", hash = "sha256:0d009c083ea859a71b76adf7c1d502e4bc170b80a8ef002da5806527b9591fac"},
+    {file = "pytest-7.4.3.tar.gz", hash = "sha256:d989d136982de4e3b29dabcc838ad581c64e8ed52c11fbe86ddebd9da0818cd5"},
 ]
 
 [package.dependencies]
@@ -2641,6 +2652,7 @@ files = [
     {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"},
     {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"},
     {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"},
+    {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"},
     {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"},
     {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"},
     {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"},
@@ -2648,8 +2660,15 @@ files = [
     {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"},
     {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"},
     {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"},
+    {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"},
     {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"},
     {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
+    {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
+    {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
+    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
+    {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
+    {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
+    {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"},
     {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"},
     {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"},
     {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"},
@@ -2666,6 +2685,7 @@ files = [
     {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"},
     {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"},
     {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"},
+    {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"},
     {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"},
     {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"},
     {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"},
@@ -2673,6 +2693,7 @@ files = [
     {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"},
     {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"},
     {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"},
+    {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"},
     {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"},
     {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"},
     {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"},
@@ -3888,4 +3909,4 @@ server = ["fastapi", "sse-starlette"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8.1"
-content-hash = "a02e9f4afeff9a7fd8cf972fd48e80fb576941aa35af43c653c7b0de1a88c691"
+content-hash = "0cca716274936ceacbe60c6d9697ddcc6714a57245c8ae18518046a1329c05cf"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langserve"
-version = "0.0.15"
+version = "0.0.17"
 description = ""
 readme = "README.md"
 authors = ["LangChain"]
@@ -16,7 +16,7 @@ fastapi = {version = ">=0.90.1", optional = true}
 sse-starlette = {version = "^1.3.0", optional = true}
 httpx-sse = {version = ">=0.3.1", optional = true}
 pydantic = "^1"
-langchain = ">=0.0.316"
+langchain = ">=0.0.322"
 
 [tool.poetry.group.dev.dependencies]
 jupyterlab = "^3.6.1"

tests/unit_tests/test_callbacks.py

@@ -0,0 +1,62 @@
+import uuid
+
+import pytest
+
+from langserve.callbacks import AsyncEventAggregatorCallback, replace_uuids
+
+
+@pytest.mark.asyncio
+async def test_event_aggregator() -> None:
+    """Test that the event aggregator is aggregating events."""
+
+    from langchain.llms import FakeListLLM
+    from langchain.prompts import ChatPromptTemplate
+
+    prompt = ChatPromptTemplate.from_template("{question}")
+    llm = FakeListLLM(responses=["hello", "world"])
+
+    chain = prompt | llm
+    callback = AsyncEventAggregatorCallback()
+    assert callback.callback_events == []
+    assert chain.invoke({"question": "hello"}, {"callbacks": [callback]}) == "hello"
+    callback_events = callback.callback_events
+    assert isinstance(callback_events, list)
+    assert len(callback_events) == 6
+    assert [event["type"] for event in callback_events] == [
+        "on_chain_start",
+        "on_chain_start",
+        "on_chain_end",
+        "on_llm_start",
+        "on_llm_end",
+        "on_chain_end",
+    ]
+
+
+def test_replace_uuids() -> None:
+    """Test replace uuids in place."""
+    uuid1 = uuid.UUID(int=1)
+    uuid2 = uuid.UUID(int=2)
+
+    events = [
+        {
+            "type": "on_llm_start",
+            "run_id": uuid1,
+            "parent_run_id": None,
+        },
+        {
+            "type": "on_llm_start",
+            "run_id": uuid1,
+            "parent_run_id": uuid2,
+        },
+    ]
+    new_events = replace_uuids(events)
+    # Assert original event is unchanged
+    assert events[0]["run_id"] == uuid1
+    assert isinstance(new_events, list)
+    assert len(new_events) == 2
+
+    assert new_events[0]["run_id"] != uuid1
+    assert new_events[1]["run_id"] != uuid2
+    assert new_events[0]["run_id"] == new_events[1]["run_id"]
+    assert new_events[0]["run_id"] != new_events[1]["parent_run_id"]
+    assert new_events[0]["parent_run_id"] is None

tests/unit_tests/test_serialization.py

@@ -1,3 +1,4 @@
+import uuid
 from typing import Any
 
 import pytest
@@ -6,13 +7,14 @@ from langchain.schema.messages import (
     HumanMessageChunk,
     SystemMessage,
 )
+from langchain.schema.output import ChatGeneration
 
 try:
     from pydantic.v1 import BaseModel
 except ImportError:
     from pydantic import BaseModel
 
-from langserve.serialization import simple_dumps, simple_loads
+from langserve.serialization import WellKnownLCSerializer, load_events
 
 
 @pytest.mark.parametrize(
@@ -39,17 +41,22 @@ from langserve.serialization import simple_dumps, simple_loads
             "numbers": [1, 2, 3],
             "boom": "Hello, world!",
         },
+        [ChatGeneration(message=HumanMessage(content="Hello"))],
     ],
 )
 def test_serialization(data: Any) -> None:
     """There and back again! :)"""
     # Test encoding
-    assert isinstance(simple_dumps(data), str)
+    lc_serializer = WellKnownLCSerializer()
+
+    assert isinstance(lc_serializer.dumps(data), str)
+    # Translate to python primitives and load back into object
+    assert lc_serializer.loadd(lc_serializer.dumpd(data)) == data
     # Test simple equality (does not include pydantic class names)
-    assert simple_loads(simple_dumps(data)) == data
+    assert lc_serializer.loads(lc_serializer.dumps(data)) == data
     # Test full representation equality (includes pydantic class names)
     assert _get_full_representation(
-        simple_loads(simple_dumps(data))
+        lc_serializer.loads(lc_serializer.dumps(data))
     ) == _get_full_representation(data)
 
 
@@ -75,3 +82,40 @@ def _get_full_representation(data: Any) -> Any:
         return data.schema()
     else:
         return data
+
+
+@pytest.mark.parametrize(
+    "data,expected",
+    [
+        ([], []),
+        (
+            [
+                {
+                    "type": "on_llm_start",
+                    "serialized": {},
+                    "prompts": [],
+                    "run_id": str(uuid.UUID(int=2)),
+                    "parent_run_id": str(uuid.UUID(int=1)),
+                    "tags": ["h"],
+                    "metadata": {},
+                    "kwargs": {},
+                }
+            ],
+            [
+                {
+                    "type": "on_llm_start",
+                    "serialized": {},
+                    "prompts": [],
+                    "run_id": uuid.UUID(int=2),
+                    "parent_run_id": uuid.UUID(int=1),
+                    "tags": ["h"],
+                    "metadata": {},
+                    "kwargs": {},
+                }
+            ],
+        ),
+    ],
+)
+def test_decode_events(data: Any, expected: Any) -> None:
+    """Test decoding events."""
+    assert load_events(data) == expected

tests/unit_tests/test_server_client.py

@@ -3,7 +3,7 @@ import asyncio
 import json
 from asyncio import AbstractEventLoop
 from contextlib import asynccontextmanager, contextmanager
-from typing import Any, Dict, Iterator, List, Optional, Sequence, Union
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 import httpx
 import pytest
@@ -22,19 +22,20 @@ from pytest_mock import MockerFixture
 from typing_extensions import TypedDict
 
 from langserve import server
+from langserve.callbacks import AsyncEventAggregatorCallback
 from langserve.client import RemoteRunnable
 from langserve.lzstring import LZString
 from langserve.server import (
     _rename_pydantic_model,
     _replace_non_alphanumeric_with_underscores,
-    add_routes,
 )
-from tests.unit_tests.utils import FakeListLLM
 
 try:
     from pydantic.v1 import BaseModel, Field
 except ImportError:
     from pydantic import BaseModel, Field
+from langserve.server import add_routes
+from tests.unit_tests.utils import FakeListLLM, FakeTracer
 
 
 @pytest.fixture(scope="session")
@@ -63,7 +64,9 @@ def app(event_loop: AbstractEventLoop) -> FastAPI:
     runnable_lambda = RunnableLambda(func=add_one_or_passthrough)
     app = FastAPI()
     try:
-        add_routes(app, runnable_lambda, config_keys=["tags"])
+        add_routes(
+            app, runnable_lambda, config_keys=["tags"], include_callback_events=True
+        )
         yield app
     finally:
         del app
@@ -156,19 +159,24 @@ def test_server(app: FastAPI) -> None:
 
     # Test invoke
     response = sync_client.post("/invoke", json={"input": 1})
-    assert response.json() == {"output": 2}
+    assert response.json()["output"] == 2
+    events = response.json()["callback_events"]
+    assert [event["type"] for event in events] == ["on_chain_start", "on_chain_end"]
 
     # Test batch
-    response = sync_client.post("/batch", json={"inputs": [1]})
-    assert response.json() == {
-        "output": [2],
-    }
+    response = sync_client.post("/batch", json={"inputs": [2, 3]})
+    assert response.json()["output"] == [3, 4]
+
+    events = response.json()["callback_events"]
+    assert [event["type"] for event in events[0]] == ["on_chain_start", "on_chain_end"]
+
+    assert [event["type"] for event in events[1]] == ["on_chain_start", "on_chain_end"]
 
     # Test schema
     input_schema = sync_client.get("/input_schema").json()
     assert isinstance(input_schema, dict)
     assert input_schema["title"] == "RunnableLambdaInput"
-
+    #
     output_schema = sync_client.get("/output_schema").json()
     assert isinstance(output_schema, dict)
     assert output_schema["title"] == "RunnableLambdaOutput"
@@ -178,11 +186,22 @@ def test_server(app: FastAPI) -> None:
     assert output_schema["title"] == "RunnableLambdaConfig"
 
     # TODO(Team): Fix test. Issue with eventloops right now when using sync client
-    ## Test stream
+    # # Test stream
     # response = sync_client.post("/stream", json={"input": 1})
     # assert response.text == "event: data\r\ndata: 2\r\n\r\nevent: end\r\n\r\n"
 
 
+def test_serve_playground(app: FastAPI) -> None:
+    """Test the server directly via HTTP requests."""
+    sync_client = TestClient(app=app)
+    response = sync_client.get("/playground/index.html")
+    assert response.status_code == 200
+    response = sync_client.get("/playground/i_do_not_exist.txt")
+    assert response.status_code == 404
+    response = sync_client.get("/playground//etc/passwd")
+    assert response.status_code == 404
+
+
 @pytest.mark.asyncio
 async def test_server_async(app: FastAPI) -> None:
     """Test the server directly via HTTP requests."""
@@ -190,13 +209,16 @@ async def test_server_async(app: FastAPI) -> None:
 
     # Test invoke
     response = await async_client.post("/invoke", json={"input": 1})
-    assert response.json() == {"output": 2}
+    assert response.json()["output"] == 2
+    events = response.json()["callback_events"]
+    assert [event["type"] for event in events] == ["on_chain_start", "on_chain_end"]
 
     # Test batch
-    response = await async_client.post("/batch", json={"inputs": [1]})
-    assert response.json() == {
-        "output": [2],
-    }
+    response = await async_client.post("/batch", json={"inputs": [1, 2]})
+    assert response.json()["output"] == [2, 3]
+    events = response.json()["callback_events"]
+    assert [event["type"] for event in events[0]] == ["on_chain_start", "on_chain_end"]
+    assert [event["type"] for event in events[1]] == ["on_chain_start", "on_chain_end"]
 
     # Test stream
     response = await async_client.post("/stream", json={"input": 1})
@@ -215,8 +237,9 @@ async def test_server_bound_async(app_for_config: FastAPI) -> None:
         json={"input": 1, "config": {"tags": ["another-one"]}},
     )
     assert response.status_code == 200
-    assert response.json() == {
-        "output": {"tags": ["another-one", "test"], "configurable": None}
+    assert response.json()["output"] == {
+        "tags": ["another-one", "test"],
+        "configurable": None,
     }
 
     # Test batch
@@ -225,9 +248,9 @@ async def test_server_bound_async(app_for_config: FastAPI) -> None:
         json={"inputs": [1], "config": {"tags": ["another-one"]}},
     )
     assert response.status_code == 200
-    assert response.json() == {
-        "output": [{"tags": ["another-one", "test"], "configurable": None}]
-    }
+    assert response.json()["output"] == [
+        {"tags": ["another-one", "test"], "configurable": None}
+    ]
 
     # Test stream
     response = await async_client.post(
@@ -248,6 +271,17 @@ def test_invoke(client: RemoteRunnable) -> None:
     # Test invocation with config
     assert client.invoke(1, config={"tags": ["test"]}) == 2
 
+    # Test tracing
+    tracer = FakeTracer()
+    assert client.invoke(1, config={"callbacks": [tracer]}) == 2
+    assert len(tracer.runs) == 1
+    # Light test to verify that we're picking up information about the server side
+    # function being invoked via a callback.
+    assert tracer.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer.runs[0].child_runs[0].extra["kwargs"]["name"] == "add_one_or_passthrough"
+    )
+
 
 def test_batch(client: RemoteRunnable) -> None:
     """Test sync batch."""
@@ -257,15 +291,66 @@ def test_batch(client: RemoteRunnable) -> None:
         HumanMessage(content="hello")
     ]
 
+    # Test callbacks
+    # Using a single tracer for both inputs
+    tracer = FakeTracer()
+    assert client.batch([1, 2], config={"callbacks": [tracer]}) == [2, 3]
+    assert len(tracer.runs) == 2
+    # Light test to verify that we're picking up information about the server side
+    # function being invoked via a callback.
+    assert tracer.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer.runs[0].child_runs[0].extra["kwargs"]["name"] == "add_one_or_passthrough"
+    )
+
+    assert tracer.runs[1].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer.runs[1].child_runs[0].extra["kwargs"]["name"] == "add_one_or_passthrough"
+    )
+
+    # Verify that each tracer gets its own run
+    tracer1 = FakeTracer()
+    tracer2 = FakeTracer()
+    assert client.batch(
+        [1, 2], config=[{"callbacks": [tracer1]}, {"callbacks": [tracer2]}]
+    ) == [2, 3]
+    assert len(tracer1.runs) == 1
+    assert len(tracer2.runs) == 1
+    # Light test to verify that we're picking up information about the server side
+    # function being invoked via a callback.
+    assert tracer1.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer1.runs[0].child_runs[0].extra["kwargs"]["name"]
+        == "add_one_or_passthrough"
+    )
+
+    assert tracer2.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer2.runs[0].child_runs[0].extra["kwargs"]["name"]
+        == "add_one_or_passthrough"
+    )
+
 
 @pytest.mark.asyncio
 async def test_ainvoke(async_client: RemoteRunnable) -> None:
     """Test async invoke."""
     assert await async_client.ainvoke(1) == 2
+
     assert await async_client.ainvoke(HumanMessage(content="hello")) == HumanMessage(
         content="hello"
     )
 
+    # Test tracing
+    tracer = FakeTracer()
+    assert await async_client.ainvoke(1, config={"callbacks": [tracer]}) == 2
+    assert len(tracer.runs) == 1
+    # Light test to verify that we're picking up information about the server side
+    # function being invoked via a callback.
+    assert tracer.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer.runs[0].child_runs[0].extra["kwargs"]["name"] == "add_one_or_passthrough"
+    )
+
 
 @pytest.mark.asyncio
 async def test_abatch(async_client: RemoteRunnable) -> None:
@@ -276,6 +361,45 @@ async def test_abatch(async_client: RemoteRunnable) -> None:
         HumanMessage(content="hello")
     ]
 
+    # Test callbacks
+    # Using a single tracer for both inputs
+    tracer = FakeTracer()
+    assert await async_client.abatch([1, 2], config={"callbacks": [tracer]}) == [2, 3]
+    assert len(tracer.runs) == 2
+    # Light test to verify that we're picking up information about the server side
+    # function being invoked via a callback.
+    assert tracer.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer.runs[0].child_runs[0].extra["kwargs"]["name"] == "add_one_or_passthrough"
+    )
+
+    assert tracer.runs[1].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer.runs[1].child_runs[0].extra["kwargs"]["name"] == "add_one_or_passthrough"
+    )
+
+    # Verify that each tracer gets its own run
+    tracer1 = FakeTracer()
+    tracer2 = FakeTracer()
+    assert await async_client.abatch(
+        [1, 2], config=[{"callbacks": [tracer1]}, {"callbacks": [tracer2]}]
+    ) == [2, 3]
+    assert len(tracer1.runs) == 1
+    assert len(tracer2.runs) == 1
+    # Light test to verify that we're picking up information about the server side
+    # function being invoked via a callback.
+    assert tracer1.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer1.runs[0].child_runs[0].extra["kwargs"]["name"]
+        == "add_one_or_passthrough"
+    )
+
+    assert tracer2.runs[0].child_runs[0].name == "RunnableLambda"
+    assert (
+        tracer2.runs[0].child_runs[0].extra["kwargs"]["name"]
+        == "add_one_or_passthrough"
+    )
+
 
 # TODO(Team): Determine how to test
 # Some issue with event loops
@@ -547,7 +671,7 @@ async def test_input_validation(
 
     config = {"tags": ["test"], "metadata": {"a": 5}}
 
-    invoke_spy_1 = mocker.spy(server_runnable, "ainvoke")
+    server_runnable_spy = mocker.spy(server_runnable, "ainvoke")
     # Verify config is handled correctly
     async with get_async_client(app, path="/add_one") as runnable1:
         # Verify that can be invoked with valid input
@@ -555,18 +679,18 @@ async def test_input_validation(
         assert await runnable1.ainvoke(1, config=config) == 2
         # Config should be ignored but default debug information
         # will still be added
-        config_seen = invoke_spy_1.call_args[1]["config"]
+        config_seen = server_runnable_spy.call_args[0][1]
         assert "metadata" in config_seen
         assert "__useragent" in config_seen["metadata"]
         assert "__langserve_version" in config_seen["metadata"]
 
-    invoke_spy_2 = mocker.spy(server_runnable2, "ainvoke")
+    server_runnable2_spy = mocker.spy(server_runnable2, "ainvoke")
     async with get_async_client(app, path="/add_one_config") as runnable2:
         # Config accepted for runnable2
         assert await runnable2.ainvoke(1, config=config) == 2
         # Config ignored
 
-        config_seen = invoke_spy_2.call_args[1]["config"]
+        config_seen = server_runnable2_spy.call_args[0][1]
         assert config_seen["tags"] == ["test"]
         assert config_seen["metadata"]["a"] == 5
         assert "__useragent" in config_seen["metadata"]
@@ -578,10 +702,12 @@ async def test_input_validation_with_lc_types(event_loop: AbstractEventLoop) ->
     """Test client side and server side exceptions."""
 
     app = FastAPI()
-    # Test with langchain objects
-    add_routes(
-        app, RunnablePassthrough(), input_type=List[HumanMessage], config_keys=["tags"]
-    )
+
+    class InputType(TypedDict):
+        messages: List[HumanMessage]
+
+    runnable = RunnablePassthrough()
+    add_routes(app, runnable, config_keys=["tags"], input_type=InputType)
     # Invoke request
     async with get_async_client(app) as passthrough_runnable:
         with pytest.raises(httpx.HTTPError):
@@ -594,15 +720,17 @@ async def test_input_validation_with_lc_types(event_loop: AbstractEventLoop) ->
             await passthrough_runnable.ainvoke([SystemMessage(content="hello")])
 
         # Valid
-        result = await passthrough_runnable.ainvoke([HumanMessage(content="hello")])
+        await passthrough_runnable.ainvoke(
+            {"messages": [HumanMessage(content="hello")]}
+        )
 
         # Valid
         result = await passthrough_runnable.ainvoke(
-            [HumanMessage(content="hello")], config={"tags": ["test"]}
+            {"messages": [HumanMessage(content="hello")]}, config={"tags": ["test"]}
         )
 
-        assert isinstance(result, list)
-        assert isinstance(result[0], HumanMessage)
+        assert isinstance(result, dict)
+        assert isinstance(result["messages"][0], HumanMessage)
 
     # Batch request
     async with get_async_client(app) as passthrough_runnable:
@@ -615,10 +743,12 @@ async def test_input_validation_with_lc_types(event_loop: AbstractEventLoop) ->
             await passthrough_runnable.abatch([[SystemMessage(content="hello")]])
 
         # valid
-        result = await passthrough_runnable.abatch([[HumanMessage(content="hello")]])
+        result = await passthrough_runnable.abatch(
+            [{"messages": [HumanMessage(content="hello")]}]
+        )
         assert isinstance(result, list)
-        assert isinstance(result[0], list)
-        assert isinstance(result[0][0], HumanMessage)
+        assert isinstance(result[0], dict)
+        assert isinstance(result[0]["messages"][0], HumanMessage)
 
 
 def test_client_close() -> None:
@@ -800,7 +930,7 @@ async def test_input_config_output_schemas(event_loop: AbstractEventLoop) -> Non
         RunnableLambda(add_two),
         path="/add_two_custom",
         input_type=float,
-        output_type=Sequence[float],
+        output_type=float,
         config_keys=["tags", "configurable"],
     )
     add_routes(app, PromptTemplate.from_template("{question}"), path="/prompt_1")
@@ -820,7 +950,7 @@ async def test_input_config_output_schemas(event_loop: AbstractEventLoop) -> Non
         assert response.json() == {"title": "RunnableLambdaInput", "type": "integer"}
 
         response = await async_client.get("/add_two_custom/input_schema")
-        assert response.json() == {"title": "Input", "type": "number"}
+        assert response.json() == {"title": "RunnableBindingInput", "type": "number"}
 
         response = await async_client.get("/prompt_1/input_schema")
         assert response.json() == {
@@ -844,11 +974,7 @@ async def test_input_config_output_schemas(event_loop: AbstractEventLoop) -> Non
         }
 
         response = await async_client.get("/add_two_custom/output_schema")
-        assert response.json() == {
-            "items": {"type": "number"},
-            "title": "Output",
-            "type": "array",
-        }
+        assert response.json() == {"title": "RunnableBindingOutput", "type": "number"}
 
         # Just verify that the schema is not empty (it's pretty long)
         # and the actual value should be tested in LangChain
@@ -916,8 +1042,7 @@ async def test_input_schema_typed_dict() -> None:
     async with AsyncClient(app=app, base_url="http://localhost:9999") as client:
         res = await client.get("/input_schema")
         assert res.json() == {
-            "title": "Input",
-            "allOf": [{"$ref": "#/definitions/InputType"}],
+            "$ref": "#/definitions/InputType",
             "definitions": {
                 "InputType": {
                     "properties": {
@@ -933,6 +1058,7 @@ async def test_input_schema_typed_dict() -> None:
                     "type": "object",
                 }
             },
+            "title": "RunnableBindingInput",
         }
 
 
@@ -973,14 +1099,30 @@ async def test_server_side_error() -> None:
 
     # Invoke request
     async with get_async_client(app, raise_app_exceptions=False) as runnable:
+        callback = AsyncEventAggregatorCallback()
         with pytest.raises(httpx.HTTPStatusError) as cm:
-            assert await runnable.ainvoke(1)
+            assert await runnable.ainvoke(1, config={"callbacks": [callback]})
         assert isinstance(cm.value, httpx.HTTPStatusError)
+        assert [event["type"] for event in callback.callback_events] == [
+            "on_chain_start",
+            "on_chain_error",
+        ]
 
+        callback1 = AsyncEventAggregatorCallback()
+        callback2 = AsyncEventAggregatorCallback()
         with pytest.raises(httpx.HTTPStatusError) as cm:
-            assert await runnable.abatch([1, 2])
+            assert await runnable.abatch(
+                [1, 2], config=[{"callbacks": [callback1]}, {"callbacks": [callback2]}]
+            )
         assert isinstance(cm.value, httpx.HTTPStatusError)
-
+        assert [event["type"] for event in callback1.callback_events] == [
+            "on_chain_start",
+            "on_chain_error",
+        ]
+        assert [event["type"] for event in callback2.callback_events] == [
+            "on_chain_start",
+            "on_chain_error",
+        ]
         # Test astream
         chunks = []
         try:
@@ -1028,3 +1170,19 @@ def test_server_side_error_sync() -> None:
             assert chunks == [1, 2]
             assert e.response.status_code == 500
             assert e.response.text == "Internal Server Error"
+
+
+def test_error_on_bad_path() -> None:
+    """Test error on bad path"""
+    app = FastAPI()
+    with pytest.raises(ValueError):
+        add_routes(app, RunnableLambda(lambda foo: "hello"), path="foo")
+    add_routes(app, RunnableLambda(lambda foo: "hello"), path="/foo")
+
+
+def test_error_on_path_collision() -> None:
+    """Test error on path collision."""
+    app = FastAPI()
+    add_routes(app, RunnableLambda(lambda foo: "hello"), path="/foo")
+    with pytest.raises(ValueError):
+        add_routes(app, RunnableLambda(lambda foo: "hello"), path="/foo")

tests/unit_tests/utils.py

@@ -1,10 +1,13 @@
-from typing import Any, List, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional
+from uuid import UUID
 
 from langchain.callbacks.manager import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
+from langchain.callbacks.tracers.base import BaseTracer
 from langchain.llms.base import LLM
+from langsmith.schemas import Run
 
 
 class FakeListLLM(LLM):
@@ -52,3 +55,42 @@ class FakeListLLM(LLM):
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
         return {"responses": self.responses}
+
+
+class FakeTracer(BaseTracer):
+    """Fake tracer that records LangChain execution.
+
+    It replaces run ids with deterministic UUIDs."""
+
+    def __init__(self) -> None:
+        """Initialize the tracer."""
+        super().__init__()
+        self.runs: List[Run] = []
+        self.uuids_map: Dict[UUID, UUID] = {}
+        self.uuids_generator = (
+            UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10_000)
+        )
+
+    def _replace_uuid(self, uuid: UUID) -> UUID:
+        """Replace a UUID with a deterministic one."""
+        if uuid not in self.uuids_map:
+            self.uuids_map[uuid] = next(self.uuids_generator)
+        return self.uuids_map[uuid]
+
+    def _copy_run(self, run: Run) -> Run:
+        """Copy a run, replacing UUIDs."""
+        return run.copy(
+            update={
+                "id": self._replace_uuid(run.id),
+                "parent_run_id": self.uuids_map[run.parent_run_id]
+                if run.parent_run_id
+                else None,
+                "child_runs": [self._copy_run(child) for child in run.child_runs],
+                "execution_order": None,
+                "child_execution_order": None,
+            }
+        )
+
+    def _persist_run(self, run: Run) -> None:
+        """Persist a run."""
+        self.runs.append(self._copy_run(run))