Version 0.0.20 (#132)

Version 0.0.20

README.md

@@ -1,4 +1,6 @@
-# LangServe 🦜️🔗 
+# LangServe 🦜️🏓
+
+🚩 We will be releasing a hosted version of LangServe for one-click deployments of LangChain applications. [Sign up here](https://airtable.com/app0hN6sd93QcKubv/shrAjst60xXa6quV2) to get on the waitlist.
 
 ## Overview
 
@@ -15,6 +17,7 @@ A javascript client is available in [LangChainJS](https://js.langchain.com/docs/
 - API docs page with JSONSchema and Swagger (insert example link)
 - Efficient `/invoke`, `/batch` and `/stream` endpoints with support for many concurrent requests on a single server
 - `/stream_log` endpoint for streaming all (or some) intermediate steps from your chain/agent
+- Playground page at `/playground` with streaming output and intermediate steps
 - Built-in (optional) tracing to [LangSmith](https://www.langchain.com/langsmith), just add your API key (see [Instructions](https://docs.smith.langchain.com/)])
 - All built with battle-tested open-source Python libraries like FastAPI, Pydantic, uvloop and asyncio.
 - Use the client SDK to call a LangServe server as if it was a Runnable running locally (or call the HTTP API directly)
@@ -24,7 +27,15 @@ 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)
 
-## LangChain CLI  🛠️
+## Hosted LangServe
+
+We will be releasing a hosted version of LangServe for one-click deployments of LangChain applications. [Sign up here](https://airtable.com/app0hN6sd93QcKubv/shrAjst60xXa6quV2) to get on the waitlist.
+
+## 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.
 
@@ -41,7 +52,6 @@ And follow the instructions...
 
 For more examples, see the [examples](./examples) directory.
 
-
 ### Server
 
 Here's a server that deploys an OpenAI chat model, an Anthropic chat model, and a chain that uses
@@ -95,6 +105,14 @@ If you've deployed the server above, you can view the generated OpenAPI docs usi
 curl localhost:8000/docs
 ```
 
+make sure to **add** the `/docs` suffix. 
+
+Below will return a 404 until you define a `@app.get("/")`
+
+```sh
+localhost:8000
+```
+
 ### Client
 
 Python SDK
@@ -116,18 +134,18 @@ joke_chain.invoke({"topic": "parrots"})
 await joke_chain.ainvoke({"topic": "parrots"})
 
 prompt = [
-    SystemMessage(content='Act like either a cat or a parrot.'), 
+    SystemMessage(content='Act like either a cat or a parrot.'),
     HumanMessage(content='Hello!')
 ]
 
 # Supports astream
 async for msg in anthropic.astream(prompt):
     print(msg, end="", flush=True)
-    
+
 prompt = ChatPromptTemplate.from_messages(
     [("system", "Tell me a long story about {topic}")]
 )
-    
+
 # Can define custom chains
 chain = prompt | RunnableMap({
     "openai": openai,
@@ -142,9 +160,11 @@ In TypeScript (requires LangChain.js version 0.0.166 or later):
 ```typescript
 import { RemoteRunnable } from "langchain/runnables/remote";
 
-const chain = new RemoteRunnable({ url: `http://localhost:8000/chain/invoke/` });
+const chain = new RemoteRunnable({
+  url: `http://localhost:8000/chain/invoke/`,
+});
 const result = await chain.invoke({
-  "topic": "cats", 
+  topic: "cats",
 });
 ```
 
@@ -171,8 +191,7 @@ curl --location --request POST 'http://localhost:8000/chain/invoke/' \
     }'
 ```
 
-
-## Endpoints 
+## Endpoints
 
 The following code:
 
@@ -195,6 +214,10 @@ adds of these endpoints to the server:
 - `GET /my_runnable/output_schema` - json schema for output of the runnable
 - `GET /my_runnable/config_schema` - json schema for config of the runnable
 
+## Playground
+
+You can find a playground page for your runnable at `/my_runnable/playground`. This exposes a simple UI to [configure](https://python.langchain.com/docs/expression_language/how_to/configure) and invoke your runnable with streaming output and intermediate steps.
+
 ## Installation
 
 For both client and server:
@@ -212,10 +235,9 @@ However, some of the input schemas for legacy chains may be incomplete/incorrect
 This can be fixed by updating the `input_schema` property of those chains in LangChain.
 If you encounter any errors, please open an issue on THIS repo, and we will work to address it.
 
-
 ## Handling Authentication
 
-If you need to add authentication to your server, 
+If you need to add authentication to your server,
 please reference FastAPI's [security documentation](https://fastapi.tiangolo.com/tutorial/security/)
 and [middleware documentation](https://fastapi.tiangolo.com/tutorial/middleware/).
 
@@ -228,3 +250,143 @@ You can deploy to GCP Cloud Run using the following command:
 ```
 gcloud run deploy [your-service-name] --source . --port 8001 --allow-unauthenticated --region us-central1 --set-env-vars=OPENAI_API_KEY=your_key
 ```
+
+## Advanced
+
+### Files
+
+LLM applications often deal with files. There are different architectures
+that can be made to implement file processing; at a high level:
+
+1. The file may be uploaded to the server via a dedicated endpoint and processed using a separate endpoint
+2. The file may be uploaded by either value (bytes of file) or reference (e.g., s3 url to file content)
+3. The processing endpoint may be blocking or non-blocking
+4. If significant processing is required, the processing may be offloaded to a dedicated process pool
+
+You should determine what is the appropriate architecture for your application.
+
+Currently, to upload files by value to a runnable, use base64 encoding for the 
+file (`multipart/form-data` is not supported yet). 
+
+Here's an [example](https://github.com/langchain-ai/langserve/tree/main/examples/file_processing) that shows
+how to use base64 encoding to send a file to a remote runnable.
+
+Remember, you can always upload files by reference (e.g., s3 url) or upload them as
+multipart/form-data to a dedicated endpoint.
+
+### Custom Input and Output Types
+
+Input and Output types are defined on all runnables.
+
+You can access them via the `input_schema` and `output_schema` properties.
+
+`LangServe` uses these types for validation and documentation.
+
+If you want to override the default inferred types, you can use the `with_types` method.
+
+Here's a toy example to illustrate the idea:
+
+```python
+from typing import Any
+
+from fastapi import FastAPI
+from langchain.schema.runnable import RunnableLambda
+
+app = FastAPI()
+
+
+def func(x: Any) -> int:
+    """Mistyped function that should accept an int but accepts anything."""
+    return x + 1
+
+
+runnable = RunnableLambda(func).with_types(
+    input_schema=int,
+)
+
+add_routes(app, runnable)
+```
+
+### Custom User Types
+
+Inherit from `CustomUserType` if you want the data to de-serialize into a 
+pydantic model rather than the equivalent dict representation.
+
+At the moment, this type only works *server* side and is used
+to specify desired *decoding* behavior. If inheriting from this type
+the server will keep the decoded type as a pydantic model instead
+of converting it into a dict.
+
+```python
+from fastapi import FastAPI
+from langchain.schema.runnable import RunnableLambda
+
+from langserve import add_routes
+from langserve.schema import CustomUserType
+
+app = FastAPI()
+
+
+class Foo(CustomUserType):
+    bar: int
+
+
+def func(foo: Foo) -> int:
+    """Sample function that expects a Foo type which is a pydantic model"""
+    assert isinstance(foo, Foo)
+    return foo.bar
+
+# Note that the input and output type are automatically inferred!
+# You do not need to specify them.
+# runnable = RunnableLambda(func).with_types( # <-- Not needed in this case
+#     input_schema=Foo,
+#     output_schema=int,
+# 
+add_routes(app, RunnableLambda(func), path="/foo")
+```
+
+### Playground Widgets
+
+The playground allows you to define custom widgets for your runnable from the backend.
+
+- A widget is specified at the field level and shipped as part of the JSON schema of the input type
+- A widget must contain a key called `type` with the value being one of a well known list of widgets
+- Other widget keys will be associated with values that describe paths in a JSON object
+
+General schema:
+
+```typescript
+type JsonPath = number | string | (number | string)[];
+type NameSpacedPath = { title: string; path: JsonPath }; // Using title to mimick json schema, but can use namespace
+type OneOfPath = { oneOf: JsonPath[] };
+
+type Widget = {
+    type: string // Some well known type (e.g., base64file, chat etc.)
+    [key: string]: JsonPath | NameSpacedPath | OneOfPath;
+};
+```
+
+
+#### File Upload Widget
+
+Allows creation of a file upload input in the UI playground for files
+that are uploaded as base64 encoded strings. Here's the full [example](https://github.com/langchain-ai/langserve/tree/main/examples/file_processing).
+
+Snippet:
+
+```python
+from pydantic import Field
+
+from langserve import CustomUserType
+
+
+# ATTENTION: Inherit from CustomUserType instead of BaseModel otherwise
+#            the server will decode it into a dict instead of a pydantic model.
+class FileProcessingRequest(CustomUserType):
+    """Request including a base64 encoded file."""
+
+    # The extra field is used to specify a widget for the playground UI.
+    file: str = Field(..., extra={"widget": {"type": "base64file"}})
+    num_chars: int = 100
+
+```

examples/chain/server.py

@@ -4,8 +4,6 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from langchain.chat_models import ChatOpenAI
 from langchain.prompts import PromptTemplate
-
-# from typing_extensions import TypedDict
 from langchain.pydantic_v1 import BaseModel
 from langchain.schema.output_parser import StrOutputParser
 from langchain.schema.runnable import ConfigurableField

examples/file_processing/client.ipynb

@@ -0,0 +1,156 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# File processing\n",
+    "\n",
+    "This client will be uploading a PDF file to the langserve server which will read the PDF and extract content from the first page."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's load the file in base64 encoding:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import base64\n",
+    "\n",
+    "with open(\"sample.pdf\", \"rb\") as f:\n",
+    "    data = f.read()\n",
+    "\n",
+    "encoded_data = base64.b64encode(data).decode(\"utf-8\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Using raw requests"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'output': 'If you’re reading this you might be using LangServe 🦜🏓!\\n\\nThis is a sample PDF!\\n\\n\\x0c',\n",
+       " 'callback_events': []}"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import requests\n",
+    "\n",
+    "requests.post(\n",
+    "    \"http://localhost:8000/pdf/invoke/\", json={\"input\": {\"file\": encoded_data}}\n",
+    ").json()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Using the SDK"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from langserve import RemoteRunnable\n",
+    "\n",
+    "runnable = RemoteRunnable(\"http://localhost:8000/pdf/\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'If you’re reading this you might be using LangServe 🦜🏓!\\n\\nThis is a sample PDF!\\n\\n\\x0c'"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "runnable.invoke({\"file\": encoded_data})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['If you’re reading this you might be using LangServe 🦜🏓!\\n\\nThis is a sample PDF!\\n\\n\\x0c',\n",
+       " 'If you’re ']"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "runnable.batch([{\"file\": encoded_data}, {\"file\": encoded_data, \"num_chars\": 10}])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

examples/file_processing/server.py

@@ -0,0 +1,61 @@
+"""Example that shows how to upload files and process files in the server.
+
+This example uses a very simple architecture for dealing with file uploads
+and processing.
+
+The main issue with this approach is that processing is done in
+the same process rather than offloaded to a process pool. A smaller
+issue is that the base64 encoding incurs an additional encoding/decoding
+overhead.
+
+This example also specifies a "base64file" widget, which will create a widget
+allowing one to upload a binary file using the langserve playground UI.
+"""
+import base64
+
+from fastapi import FastAPI
+from langchain.document_loaders.blob_loaders import Blob
+from langchain.document_loaders.parsers.pdf import PDFMinerParser
+from langchain.schema.runnable import RunnableLambda
+from pydantic import Field
+
+from langserve import CustomUserType, add_routes
+
+app = FastAPI(
+    title="LangChain Server",
+    version="1.0",
+    description="Spin up a simple api server using Langchain's Runnable interfaces",
+)
+
+
+# ATTENTION: Inherit from CustomUserType instead of BaseModel otherwise
+#            the server will decode it into a dict instead of a pydantic model.
+class FileProcessingRequest(CustomUserType):
+    """Request including a base64 encoded file."""
+
+    # The extra field is used to specify a widget for the playground UI.
+    file: str = Field(..., extra={"widget": {"type": "base64file"}})
+    num_chars: int = 100
+
+
+def _process_file(request: FileProcessingRequest) -> str:
+    """Extract the text from the first page of the PDF."""
+    content = base64.b64decode(request.file.encode("utf-8"))
+    blob = Blob(data=content)
+    documents = list(PDFMinerParser().lazy_parse(blob))
+    content = documents[0].page_content
+    return content[: request.num_chars]
+
+
+add_routes(
+    app,
+    RunnableLambda(_process_file).with_types(input_type=FileProcessingRequest),
+    config_keys=["configurable"],
+    path="/pdf",
+)
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="localhost", port=8000)

examples/llm/client.ipynb

@@ -229,7 +229,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain.schema.runnable import PutLocalVar, GetLocalVar"
+    "from langchain.schema.runnable import RunnablePassthrough"
    ]
   },
   {
@@ -266,13 +266,8 @@
     ")\n",
     "\n",
     "\n",
-    "chain = (\n",
-    "    comedian_chain\n",
-    "    | PutLocalVar(\"joke\")\n",
-    "    | {\"joke\": GetLocalVar(\"joke\")}\n",
-    "    | joke_classifier_chain\n",
-    "    | PutLocalVar(\"classification\")\n",
-    "    | {\"joke\": GetLocalVar(\"joke\"), \"classification\": GetLocalVar(\"classification\")}\n",
+    "chain = {\"joke\": comedian_chain} | RunnablePassthrough.assign(\n",
+    "    classification=joke_classifier_chain\n",
     ")"
    ]
   },

examples/widgets/server.py

@@ -0,0 +1,80 @@
+import base64
+from typing import Any, Dict, List, Tuple
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from langchain.document_loaders.blob_loaders import Blob
+from langchain.document_loaders.parsers.pdf import PDFMinerParser
+from langchain.schema.runnable import RunnableLambda
+from pydantic import BaseModel, Field
+
+from langserve.server import add_routes
+
+app = FastAPI(
+    title="LangChain Server",
+    version="1.0",
+    description="Spin up a simple api server using Langchain's Runnable interfaces",
+)
+
+
+# Set all CORS enabled origins
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+    expose_headers=["*"],
+)
+
+
+class ChatHistory(BaseModel):
+    chat_history: List[Tuple[str, str]] = Field(
+        ...,
+        examples=[[("a", "aa")]],
+        extra={"widget": {"type": "chat", "input": "question", "output": "answer"}},
+    )
+    question: str
+
+
+class FileProcessingRequest(BaseModel):
+    file: bytes = Field(..., extra={"widget": {"type": "base64file"}})
+    num_chars: int = 100
+
+
+def chat_with_bot(input: Dict[str, Any]) -> Dict[str, Any]:
+    """Bot that repeats the question twice."""
+    return {
+        "answer": input["question"] * 2,
+        "woof": "its so bad to woof, meow is better",
+    }
+
+
+def process_file(input: Dict[str, Any]) -> str:
+    """Extract the text from the first page of the PDF."""
+    content = base64.decodebytes(input["file"])
+    blob = Blob(data=content)
+    documents = list(PDFMinerParser().lazy_parse(blob))
+    content = documents[0].page_content
+    return content[: input["num_chars"]]
+
+
+add_routes(
+    app,
+    RunnableLambda(chat_with_bot).with_types(input_type=ChatHistory),
+    config_keys=["configurable"],
+    path="/chat",
+)
+
+add_routes(
+    app,
+    RunnableLambda(process_file).with_types(input_type=FileProcessingRequest),
+    config_keys=["configurable"],
+    path="/pdf",
+)
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="localhost", port=8000)

langserve/__init__.py

@@ -1,7 +1,12 @@
-"""Main entrypoint into package."""
+"""Main entrypoint into package.
+
+This is the ONLY public interface into the package. All other modules are
+to be considered private and subject to change without notice.
+"""
 
 from langserve.client import RemoteRunnable
+from langserve.schema import CustomUserType
 from langserve.server import add_routes
 from langserve.version import __version__
 
-__all__ = ["RemoteRunnable", "add_routes", "__version__"]
+__all__ = ["RemoteRunnable", "add_routes", "__version__", "CustomUserType"]

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,8 +1,12 @@
 from __future__ import annotations
 
 import asyncio
+import copy
+import json
+import logging
 import weakref
 from concurrent.futures import ThreadPoolExecutor
+from functools import lru_cache
 from typing import (
     Any,
     AsyncIterator,
@@ -11,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
@@ -28,7 +37,14 @@ 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__)
 
 
 def _without_callbacks(config: Optional[RunnableConfig]) -> RunnableConfig:
@@ -37,6 +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 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.
 
@@ -58,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,
         )
 
@@ -94,6 +139,62 @@ def _close_clients(sync_client: httpx.Client, async_client: httpx.AsyncClient) -
         asyncio.run(async_client.aclose())
 
 
+def _raise_exception_from_data(data: str, request: httpx.Request) -> None:
+    """Raise an httpx exception from the given error event data."""
+    try:
+        decoded_data = json.loads(data)
+    except json.JSONDecodeError:
+        raise httpx.HTTPStatusError(
+            message="invalid json in error event sent from server",
+            request=_sanitize_request(request),
+            response=httpx.Response(status_code=500, text=data),
+        )
+    raise httpx.HTTPStatusError(
+        message=decoded_data["message"],
+        request=_sanitize_request(request),
+        response=httpx.Response(
+            status_code=decoded_data["status_code"],
+            text=decoded_data["message"],
+        ),
+    )
+
+
+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.
 
@@ -103,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__(
@@ -118,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.
 
@@ -130,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
@@ -157,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
@@ -181,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
@@ -204,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,
@@ -224,13 +347,28 @@ 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 _enforce_trailing_slash(self, url: str) -> str:
+        if url.endswith("/"):
+            return url
+        return url + "/"
 
     def batch(
         self,
@@ -245,6 +383,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,
@@ -266,13 +405,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,
@@ -303,15 +456,15 @@ 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,
         }
-        endpoint = urljoin(self.url, "stream")
+        endpoint = urljoin(self._enforce_trailing_slash(self.url), "stream")
 
         try:
             from httpx_sse import connect_sse
@@ -327,17 +480,26 @@ 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:
                             final_output += chunk
                         else:
                             final_output = chunk
+                    elif sse.event == "error":
+                        # This can only be a server side error
+                        _raise_exception_from_data(
+                            sse.data, httpx.Request(method="POST", url=endpoint)
+                        )
                     elif sse.event == "end":
                         break
                     else:
-                        raise NotImplementedError(f"Unknown event {sse.event}")
+                        logger.error(
+                            f"Encountered an unsupported event type: `{sse.event}`. "
+                            f"Try upgrading the remote client to the latest version."
+                            f"Ignoring events of type `{sse.event}`."
+                        )
         except BaseException as e:
             run_manager.on_chain_error(e)
             raise
@@ -357,15 +519,15 @@ 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,
         }
-        endpoint = urljoin(self.url, "stream")
+        endpoint = urljoin(self._enforce_trailing_slash(self.url), "stream")
 
         try:
             from httpx_sse import aconnect_sse
@@ -378,18 +540,27 @@ 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:
                             final_output += chunk
                         else:
                             final_output = chunk
+
+                    elif sse.event == "error":
+                        # This can only be a server side error
+                        _raise_exception_from_data(
+                            sse.data, httpx.Request(method="POST", url=endpoint)
+                        )
                     elif sse.event == "end":
                         break
-
                     else:
-                        raise NotImplementedError(f"Unknown event {sse.event}")
+                        logger.error(
+                            f"Encountered an unsupported event type: `{sse.event}`. "
+                            f"Try upgrading the remote client to the latest version."
+                            f"Ignoring events of type `{sse.event}`."
+                        )
         except BaseException as e:
             await run_manager.on_chain_error(e)
             raise
@@ -428,11 +599,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,
@@ -456,19 +627,27 @@ 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
-
                         if final_output:
                             final_output += chunk
                         else:
                             final_output = chunk
+                    elif sse.event == "error":
+                        # This can only be a server side error
+                        _raise_exception_from_data(
+                            sse.data, httpx.Request(method="POST", url=endpoint)
+                        )
                     elif sse.event == "end":
                         break
                     else:
-                        raise NotImplementedError(f"Unknown event {sse.event}")
+                        _log_error_message_once(
+                            f"Encountered an unsupported event type: `{sse.event}`. "
+                            f"Try upgrading the remote client to the latest version."
+                            f"Ignoring events of type `{sse.event}`."
+                        )
         except BaseException as e:
             await run_manager.on_chain_error(e)
             raise

langserve/packages.py

@@ -0,0 +1,104 @@
+import importlib
+import logging
+from pathlib import Path
+from typing import Generator, TypedDict, Union
+
+from fastapi import APIRouter, FastAPI
+from tomli import load
+
+from langserve.server import add_routes
+
+
+class LangServeExport(TypedDict):
+    """
+    Fields from pyproject.toml that are relevant to LangServe
+
+    Attributes:
+        module: The module to import from, tool.langserve.export_module
+        attr: The attribute to import from the module, tool.langserve.export_attr
+        package_name: The name of the package, tool.poetry.name
+    """
+
+    module: str
+    attr: str
+    package_name: str
+
+
+def get_langserve_export(filepath: Path) -> LangServeExport:
+    with open(filepath, "rb") as f:
+        data = load(f)
+    try:
+        module = data["tool"]["langserve"]["export_module"]
+        attr = data["tool"]["langserve"]["export_attr"]
+        package_name = data["tool"]["poetry"]["name"]
+    except KeyError as e:
+        raise KeyError("Invalid LangServe PyProject.toml") from e
+    return LangServeExport(module=module, attr=attr, package_name=package_name)
+
+
+EXCLUDE_PATHS = set(["__pycache__", ".venv", ".git", ".github"])
+
+
+def _include_path(path: Path) -> bool:
+    """
+    Skip paths that are in EXCLUDE_PATHS or start with an underscore.
+    """
+    for part in path.parts:
+        if part in EXCLUDE_PATHS:
+            return False
+        if part.startswith("_"):
+            return False
+    return True
+
+
+def list_packages(path: str = "../packages") -> Generator[Path, None, None]:
+    """
+    Yields Path objects for each folder that contains a pyproject.toml file within a
+    path. Use this to find packages to add to the server.
+
+    See `add_package_routes` below for an example of how to use this.
+    """
+    # traverse path for routes to host (any directory holding a pyproject.toml file)
+    package_root = Path(path)
+    for pyproject_path in package_root.glob("**/pyproject.toml"):
+        if not _include_path(pyproject_path):
+            continue
+        yield pyproject_path.parent
+
+
+def add_package_route(
+    app: Union[FastAPI, APIRouter], package_path: Path, mount_path: str
+) -> None:
+    try:
+        pyproject_path = package_path / "pyproject.toml"
+
+        # get langserve export
+        package = get_langserve_export(pyproject_path)
+        package_name = package["package_name"]
+        # import module
+        mod = importlib.import_module(package["module"])
+    except KeyError:
+        logging.warning(
+            f"Skipping {package_path} because it is not a valid LangServe "
+            "package (see pyproject.toml)"
+        )
+        return
+    except ImportError as e:
+        logging.warning(f"Error: {e}")
+        logging.warning(f"Try fixing with `poetry add --editable {package_path}`")
+        logging.warning(
+            "To remove packages, use `poe` instead of `poetry`: "
+            f"`poe remove {package_name}`"
+        )
+        return
+    # get attr
+    chain = getattr(mod, package["attr"])
+    add_routes(app, chain, path=mount_path)
+
+
+def add_package_routes(app: Union[FastAPI, APIRouter], path: str = "packages") -> None:
+    # traverse path for routes to host (any directory holding a pyproject.toml file)
+    for package_path in list_packages(path):
+        mount_path_relative = package_path.relative_to(Path(path))
+        mount_path = f"/{mount_path_relative}"
+        add_package_route(app, package_path, mount_path)

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-6c8f83bb.js"></script>
-    <link rel="stylesheet" href="/____LANGSERVE_BASE_URL/assets/index-5008c8a8.css">
+    <script type="module" crossorigin src="/____LANGSERVE_BASE_URL/assets/index-03e15490.js"></script>
+    <link rel="stylesheet" href="/____LANGSERVE_BASE_URL/assets/index-53e98e83.css">
   </head>
   <body>
     <div id="root"></div>

langserve/playground/package.json

@@ -23,12 +23,12 @@
     "clsx": "^2.0.0",
     "dayjs": "^1.11.10",
     "fast-json-patch": "^3.1.1",
-    "json-schema-defaults": "^0.4.0",
     "lodash": "^4.17.21",
     "lz-string": "^1.5.0",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "tailwind-merge": "^1.14.0",
+    "use-debounce": "^9.0.4",
     "vaul": "^0.7.3"
   },
   "devDependencies": {

langserve/playground/src/App.tsx

@@ -1,13 +1,11 @@
 import "./App.css";
 
-import React, { useEffect, useRef, useState } from "react";
-import defaults from "json-schema-defaults";
+import { useEffect, useMemo, useRef, useState } from "react";
+import defaults from "./utils/defaults";
 import { JsonForms } from "@jsonforms/react";
 import {
   materialAllOfControlTester,
   MaterialAllOfRenderer,
-  materialAnyOfControlTester,
-  MaterialAnyOfRenderer,
   MaterialObjectRenderer,
   materialOneOfControlTester,
   MaterialOneOfRenderer,
@@ -17,7 +15,6 @@ import utc from "dayjs/plugin/utc";
 import relativeDate from "dayjs/plugin/relativeTime";
 import SendIcon from "./assets/SendIcon.svg?react";
 import ShareIcon from "./assets/ShareIcon.svg?react";
-import ChevronRight from "./assets/ChevronRight.svg?react";
 import { compressToEncodedURIComponent } from "lz-string";
 
 import {
@@ -42,9 +39,8 @@ import {
   vanillaRenderers,
   InputControl,
 } from "@jsonforms/vanilla-renderers";
-
 import { useSchemas } from "./useSchemas";
-import { RunState, useStreamLog } from "./useStreamLog";
+import { useStreamLog } from "./useStreamLog";
 import {
   JsonFormsCore,
   RankedTester,
@@ -59,18 +55,29 @@ import CustomArrayControlRenderer, {
 } from "./components/CustomArrayControlRenderer";
 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";
+import {
+  fileBase64Tester,
+  FileBase64ControlRenderer,
+} from "./components/FileBase64Tester";
+import { IntermediateSteps } from "./components/IntermediateSteps";
+import { StreamOutput } from "./components/StreamOutput";
+import {
+  customAnyOfTester,
+  CustomAnyOfRenderer,
+} from "./components/CustomAnyOfRenderer";
 
 dayjs.extend(relativeDate);
 dayjs.extend(utc);
 
-function str(o: unknown): React.ReactNode {
-  return typeof o === "object"
-    ? JSON.stringify(o, null, 2)
-    : (o as React.ReactNode);
-}
-
 const isObjectWithPropertiesControl = rankWith(
   2,
   and(
@@ -85,26 +92,33 @@ const isObjectWithPropertiesControl = rankWith(
 const isObject = rankWith(1, and(uiTypeIs("Control"), schemaTypeIs("object")));
 const isElse = rankWith(1, and(uiTypeIs("Control")));
 
-const renderers = [
+export const renderers = [
   ...vanillaRenderers,
 
   // use material renderers to handle objects and json schema references
   // they should yield the rendering to simpler cells
   { tester: isObjectWithPropertiesControl, renderer: MaterialObjectRenderer },
   { tester: materialAllOfControlTester, renderer: MaterialAllOfRenderer },
-  { tester: materialAnyOfControlTester, renderer: MaterialAnyOfRenderer },
   { tester: materialOneOfControlTester, renderer: MaterialOneOfRenderer },
 
+  { tester: customAnyOfTester, renderer: CustomAnyOfRenderer },
+
   // custom renderers
   { tester: materialArrayControlTester, renderer: CustomArrayControlRenderer },
   { tester: isObject, renderer: InputControl },
+  { tester: chatMessagesTester, renderer: ChatMessagesControlRenderer },
+  {
+    tester: chatMessagesTupleTester,
+    renderer: ChatMessageTuplesControlRenderer,
+  },
+  { tester: fileBase64Tester, renderer: FileBase64ControlRenderer },
 ];
 
 const nestedArrayControlTester: RankedTester = rankWith(1, (_, jsonSchema) => {
   return jsonSchema.type === "array";
 });
 
-const cells = [
+export const cells = [
   { tester: booleanCellTester, cell: BooleanCell },
   { tester: dateCellTester, cell: DateCell },
   { tester: dateTimeCellTester, cell: DateTimeCell },
@@ -119,41 +133,6 @@ const cells = [
   { tester: isElse, cell: JsonTextAreaCell },
 ];
 
-function IntermediateSteps(props: { latest: RunState }) {
-  const [expanded, setExpanded] = useState(false);
-  return (
-    <div className="flex flex-col border border-divider-700 rounded-2xl bg-background">
-      <button
-        className="font-medium text-left p-4 flex items-center justify-between"
-        onClick={() => setExpanded((open) => !open)}
-      >
-        <span>Intermediate steps</span>
-        <ChevronRight
-          className={cn("transition-all", expanded && "rotate-90")}
-        />
-      </button>
-      {expanded && (
-        <div className="flex flex-col gap-5 p-4 pt-0 divide-solid divide-y divide-divider-700 rounded-b-xl">
-          {Object.values(props.latest.logs).map((log) => (
-            <div
-              className="gap-3 flex-col min-w-0 flex bg-background pt-3 first-of-type:pt-0"
-              key={log.id}
-            >
-              <div className="flex items-center justify-between">
-                <strong className="text-sm font-medium">{log.name}</strong>
-                <p className="text-sm">{dayjs.utc(log.start_time).fromNow()}</p>
-              </div>
-              <pre className="break-words whitespace-pre-wrap min-w-0 text-sm bg-ls-gray-400 rounded-lg p-3">
-                {str(log.final_output) ?? "..."}
-              </pre>
-            </div>
-          ))}
-        </div>
-      )}
-    </div>
-  );
-}
-
 function App() {
   const [isIframe] = useState(() => window.self !== window.top);
 
@@ -164,14 +143,14 @@ function App() {
 
   // store form state
   const [configData, setConfigData] = useState<
-    Pick<JsonFormsCore, "data" | "errors">
-  >({ data: {}, errors: [] });
+    Pick<JsonFormsCore, "data" | "errors"> & { defaults: boolean }
+  >({ data: {}, errors: [], defaults: true });
 
   const [inputData, setInputData] = useState<
     Pick<JsonFormsCore, "data" | "errors">
   >({ data: null, errors: [] });
   // fetch input and config schemas from the server
-  const schemas = useSchemas();
+  const schemas = useSchemas(configData);
   // apply defaults defined in each schema
   useEffect(() => {
     if (schemas.config) {
@@ -182,8 +161,10 @@ function App() {
           initConfigData.current ??
           defaults(schemas.config),
         errors: [],
+        defaults: true,
       });
-      setInputData({ data: null, errors: [] });
+
+      setInputData({ data: defaults(schemas.input), errors: [] });
     }
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [schemas.config]);
@@ -204,7 +185,11 @@ function App() {
               const value: { config: JsonFormsCore["data"] } = message.value;
               if (Object.keys(value.config).length > 0) {
                 initConfigData.current = value.config;
-                setConfigData({ data: value.config, errors: [] });
+                setConfigData({
+                  data: value.config,
+                  errors: [],
+                  defaults: false,
+                });
                 break;
               }
             }
@@ -217,44 +202,101 @@ function App() {
     return () => window.removeEventListener("message", listener);
   }, []);
 
+  const isInputResetable = useMemo(() => {
+    if (!schemas.input) return false;
+    return (
+      JSON.stringify(defaults(schemas.input)) !== JSON.stringify(inputData.data)
+    );
+  }, [schemas.input, inputData.data]);
+
+  function onSubmit() {
+    if (
+      !stopStream &&
+      (!!inputData.errors?.length || !!configData.errors?.length)
+    ) {
+      return;
+    }
+
+    if (stopStream) {
+      stopStream();
+    } else {
+      startStream(inputData.data, configData.data);
+    }
+  }
+
+  const submitRef = useRef<(() => void) | null>(null);
+  submitRef.current = onSubmit;
+
+  useEffect(() => {
+    window.addEventListener("keydown", (e) => {
+      if (e.key === "Enter" && (e.metaKey || e.ctrlKey)) {
+        e.preventDefault();
+        submitRef.current?.();
+      }
+    });
+  }, []);
+
   return schemas.config && schemas.input ? (
     <div className="flex items-center flex-col text-ls-black bg-gradient-to-b from-[#F9FAFB] to-[#EFF8FF] min-h-[100dvh] dark:from-[#0C111C] dark:to-[#0C111C]">
       <div className="flex flex-col flex-grow gap-4 px-4 pt-6 max-w-[800px] w-full">
         <h1 className="text-2xl text-left">
           <strong>🦜 LangServe</strong> Playground
         </h1>
-        <div className="flex flex-col gap-3">
-          {!isIframe && <h2 className="text-xl font-semibold">Configure</h2>}
+        {Object.keys(schemas.config).length > 0 && (
+          <div className="flex flex-col gap-3 [&:has(.content>.vertical-layout:first-child:last-child:empty)]:hidden">
+            {!isIframe && <h2 className="text-xl font-semibold">Configure</h2>}
 
-          <JsonForms
-            schema={schemas.config}
-            data={configData.data}
-            renderers={renderers}
-            cells={cells}
-            onChange={({ data, errors }) =>
-              data ? setConfigData({ data, errors }) : undefined
-            }
-          />
-          {!!configData.errors?.length && configData.data && (
-            <div className="bg-background rounded-xl">
-              <div className="bg-red-500/10 text-red-700 dark:text-red-300 rounded-xl p-3">
-                <strong className="font-bold">Validation Errors</strong>
-                <ul className="list-disc pl-5">
-                  {configData.errors?.map((e, i) => (
-                    <li key={i}>{e.message}</li>
-                  ))}
-                </ul>
-              </div>
+            <div className="content flex flex-col gap-3">
+              <JsonForms
+                schema={schemas.config}
+                data={configData.data}
+                renderers={renderers}
+                cells={cells}
+                onChange={({ data, errors }) =>
+                  data
+                    ? setConfigData({ data, errors, defaults: false })
+                    : undefined
+                }
+              />
+
+              {!!configData.errors?.length && configData.data && (
+                <div className="bg-background rounded-xl">
+                  <div className="bg-red-500/10 text-red-700 dark:text-red-300 rounded-xl p-3">
+                    <strong className="font-bold">Validation Errors</strong>
+                    <ul className="list-disc pl-5">
+                      {configData.errors?.map((e, i) => (
+                        <li key={i}>{e.message}</li>
+                      ))}
+                    </ul>
+                  </div>
+                </div>
+              )}
             </div>
-          )}
-        </div>
+          </div>
+        )}
 
         {!isIframe && (
           <div className="flex flex-col gap-3">
             <h2 className="text-xl font-semibold">Try it</h2>
 
             <div className="p-4 border border-divider-700 flex flex-col gap-3 rounded-2xl bg-background">
-              <h3 className="font-medium">Inputs</h3>
+              <div className="flex items-center justify-between">
+                <h3 className="font-medium">Inputs</h3>
+                {isInputResetable && (
+                  <button
+                    type="button"
+                    className="text-sm px-1 -mr-1 py-0.5 rounded-md hover:bg-divider-500/50 active:bg-divider-500 text-ls-gray-100"
+                    onClick={() =>
+                      setInputData({
+                        data: defaults(schemas.input),
+                        errors: [],
+                      })
+                    }
+                  >
+                    Reset
+                  </button>
+                )}
+              </div>
 
               <JsonForms
                 schema={schemas.input}
@@ -279,7 +321,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>
@@ -338,18 +380,35 @@ function App() {
               <button
                 type="button"
                 className="px-4 py-3 gap-3 font-medium border border-transparent rounded-full flex items-center justify-center bg-blue-500 hover:bg-blue-600 active:bg-blue-700 disabled:opacity-50 transition-colors"
-                onClick={() => {
-                  stopStream
-                    ? stopStream()
-                    : startStream(inputData.data, configData.data);
-                }}
+                onClick={onSubmit}
                 disabled={
                   !stopStream &&
                   (!!inputData.errors?.length || !!configData.errors?.length)
                 }
               >
                 {stopStream ? (
-                  <span className="text-white">Stop</span>
+                  <>
+                    <div role="status">
+                      <svg
+                        aria-hidden="true"
+                        className="w-5 h-5 animate-spin text-white fill-ls-blue"
+                        viewBox="0 0 100 101"
+                        fill="none"
+                        xmlns="http://www.w3.org/2000/svg"
+                      >
+                        <path
+                          d="M100 50.5908C100 78.2051 77.6142 100.591 50 100.591C22.3858 100.591 0 78.2051 0 50.5908C0 22.9766 22.3858 0.59082 50 0.59082C77.6142 0.59082 100 22.9766 100 50.5908ZM9.08144 50.5908C9.08144 73.1895 27.4013 91.5094 50 91.5094C72.5987 91.5094 90.9186 73.1895 90.9186 50.5908C90.9186 27.9921 72.5987 9.67226 50 9.67226C27.4013 9.67226 9.08144 27.9921 9.08144 50.5908Z"
+                          fill="currentColor"
+                        />
+                        <path
+                          d="M93.9676 39.0409C96.393 38.4038 97.8624 35.9116 97.0079 33.5539C95.2932 28.8227 92.871 24.3692 89.8167 20.348C85.8452 15.1192 80.8826 10.7238 75.2124 7.41289C69.5422 4.10194 63.2754 1.94025 56.7698 1.05124C51.7666 0.367541 46.6976 0.446843 41.7345 1.27873C39.2613 1.69328 37.813 4.19778 38.4501 6.62326C39.0873 9.04874 41.5694 10.4717 44.0505 10.1071C47.8511 9.54855 51.7191 9.52689 55.5402 10.0491C60.8642 10.7766 65.9928 12.5457 70.6331 15.2552C75.2735 17.9648 79.3347 21.5619 82.5849 25.841C84.9175 28.9121 86.7997 32.2913 88.1811 35.8758C89.083 38.2158 91.5421 39.6781 93.9676 39.0409Z"
+                          fill="currentFill"
+                        />
+                      </svg>
+                      <span className="sr-only">Loading...</span>
+                    </div>
+                    <span className="text-white">Stop</span>
+                  </>
                 ) : (
                   <>
                     <SendIcon className="flex-shrink-0" />

langserve/playground/src/components/ChatMessageTuplesControlRenderer.tsx

@@ -0,0 +1,115 @@
+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 { isJsonSchemaExtra } from "../utils/schema";
+
+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 (!isJsonSchemaExtra(schema) || schema.extra.widget.type !== "chat") {
+        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")
+        );
+      }
+
+      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 || "Messages"}
+          </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,162 @@
+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";
+
+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 || "Messages"}
+          </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/CustomAnyOfRenderer.tsx

@@ -0,0 +1,37 @@
+import { JsonFormsDispatch, withJsonFormsAnyOfProps } from "@jsonforms/react";
+import {
+  rankWith,
+  createCombinatorRenderInfos,
+  JsonSchema,
+  isAnyOfControl,
+} from "@jsonforms/core";
+import { renderers, cells } from "../App";
+
+export const CustomAnyOfRenderer = withJsonFormsAnyOfProps((props) => {
+  const anyOfRenderInfos = createCombinatorRenderInfos(
+    (props.schema as JsonSchema).anyOf!,
+    props.rootSchema,
+    "anyOf",
+    props.uischema,
+    props.path,
+    props.uischemas
+  );
+
+  // just assume the last type is the selected one
+  // for `anyOf` caused by passing inputs from LLMs/Chat Models
+  // this will result in showing the Message renderer
+  const selectedIndex = anyOfRenderInfos.length - 1;
+  const selectedAnyOfRenderInfo = anyOfRenderInfos[selectedIndex];
+
+  return (
+    <JsonFormsDispatch
+      schema={selectedAnyOfRenderInfo.schema}
+      uischema={selectedAnyOfRenderInfo.uischema}
+      path={props.path}
+      renderers={renderers}
+      cells={cells}
+    />
+  );
+});
+
+export const customAnyOfTester = rankWith(3, isAnyOfControl);

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/components/FileBase64Tester.tsx

@@ -0,0 +1,44 @@
+import { ChangeEvent } from "react";
+import { withJsonFormsControlProps } from "@jsonforms/react";
+import { rankWith, and, schemaMatches, isControl } from "@jsonforms/core";
+import { isJsonSchemaExtra } from "../utils/schema";
+
+export const fileBase64Tester = rankWith(
+  12,
+  and(
+    isControl,
+    schemaMatches((schema) => {
+      if (!isJsonSchemaExtra(schema)) return false;
+      return schema.extra.widget.type === "base64file";
+    })
+  )
+);
+
+export const FileBase64ControlRenderer = withJsonFormsControlProps((props) => {
+  const handleFileUpload = (event: ChangeEvent<HTMLInputElement>) => {
+    const file = event.target.files?.[0];
+    if (!file) return;
+
+    const reader = new FileReader();
+
+    reader.onload = () => {
+      const base64String = reader.result as string | null;
+      if (base64String != null) {
+        const prefix = base64String.indexOf("base64,") + "base64,".length;
+        props.handleChange(props.path, base64String.slice(prefix));
+      }
+    };
+
+    reader.readAsDataURL(file);
+  };
+
+  return (
+    <div className="control">
+      <label className="text-xs uppercase font-semibold text-ls-gray-100">
+        {props.label}
+      </label>
+
+      <input type="file" onChange={handleFileUpload} />
+    </div>
+  );
+});

langserve/playground/src/components/IntermediateSteps.tsx

@@ -0,0 +1,41 @@
+import { useState } from "react";
+import dayjs from "dayjs";
+import ChevronRight from "../assets/ChevronRight.svg?react";
+import { RunState } from "../useStreamLog";
+import { cn } from "../utils/cn";
+import { str } from "../utils/str";
+
+export function IntermediateSteps(props: { latest: RunState }) {
+  const [expanded, setExpanded] = useState(false);
+  return (
+    <div className="flex flex-col border border-divider-700 rounded-2xl bg-background">
+      <button
+        className="font-medium text-left p-4 flex items-center justify-between"
+        onClick={() => setExpanded((open) => !open)}
+      >
+        <span>Intermediate steps</span>
+        <ChevronRight
+          className={cn("transition-all", expanded && "rotate-90")}
+        />
+      </button>
+      {expanded && (
+        <div className="flex flex-col gap-5 p-4 pt-0 divide-solid divide-y divide-divider-700 rounded-b-xl">
+          {Object.values(props.latest.logs).map((log) => (
+            <div
+              className="gap-3 flex-col min-w-0 flex bg-background pt-3 first-of-type:pt-0"
+              key={log.id}
+            >
+              <div className="flex items-center justify-between">
+                <strong className="text-sm font-medium">{log.name}</strong>
+                <p className="text-sm">{dayjs.utc(log.start_time).fromNow()}</p>
+              </div>
+              <pre className="break-words whitespace-pre-wrap min-w-0 text-sm bg-ls-gray-400 rounded-lg p-3">
+                {str(log.final_output) ?? "..."}
+              </pre>
+            </div>
+          ))}
+        </div>
+      )}
+    </div>
+  );
+}

langserve/playground/src/components/StreamOutput.tsx

@@ -0,0 +1,108 @@
+import { str } from "../utils/str";
+
+// 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));
+}
+
+export 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
+    );
+
+    const functionCall = concat?.additional_kwargs?.function_call;
+    return (
+      concat?.content ||
+      (!!functionCall && JSON.stringify(functionCall, null, 2)) ||
+      "..."
+    );
+  }
+
+  return props.streamed.map(str).join("") || "...";
+}

langserve/playground/src/types/json-schema-defaults.d.ts

@@ -1,5 +0,0 @@
-declare module "json-schema-defaults" {
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  function defaults(schema: any): any;
-  export = defaults;
-}

langserve/playground/src/useSchemas.tsx

@@ -1,6 +1,9 @@
 import { useEffect, useState } from "react";
 import { resolveApiUrl } from "./utils/url";
 import { simplifySchema } from "./utils/simplifySchema";
+import { JsonFormsCore } from "@jsonforms/core";
+import { compressToEncodedURIComponent } from "lz-string";
+import { useDebounce } from "use-debounce";
 
 declare global {
   interface Window {
@@ -11,8 +14,15 @@ declare global {
   }
 }
 
-export function useSchemas() {
-  const [schemas, setSchemas] = useState({
+export function useSchemas(
+  configData: Pick<JsonFormsCore, "data" | "errors"> & { defaults: boolean }
+) {
+  const [schemas, setSchemas] = useState<{
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    config: null | any;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    input: null | any;
+  }>({
     config: null,
     input: null,
   });
@@ -44,5 +54,23 @@ export function useSchemas() {
     save();
   }, []);
 
+  const [debouncedConfigData] = useDebounce(configData, 500);
+
+  useEffect(() => {
+    if (!debouncedConfigData.defaults) {
+      fetch(
+        resolveApiUrl(
+          `/c/${compressToEncodedURIComponent(
+            JSON.stringify(debouncedConfigData.data)
+          )}/input_schema`
+        )
+      )
+        .then((r) => r.json())
+        .then(simplifySchema)
+        .then((input) => setSchemas((current) => ({ ...current, input })))
+        .catch(() => {}); // ignore errors, eg. due to incomplete config
+    }
+  }, [debouncedConfigData]);
+
   return schemas;
 }

langserve/playground/src/utils/defaults.ts

@@ -0,0 +1,220 @@
+// (c) 2015 Chute Corporation. Released under the terms of the MIT License.
+// Modified to use TypeScript and handle edge cases with tuples
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+/* eslint-disable no-prototype-builtins */
+
+"use strict";
+
+/**
+ * check whether item is plain object
+ * @param {*} item
+ * @return {Boolean}
+ */
+const isObject = (item: unknown): item is Record<string, unknown> => {
+  return (
+    typeof item === "object" &&
+    item !== null &&
+    item.toString() === {}.toString()
+  );
+};
+
+/**
+ * deep JSON object clone
+ *
+ * @param {Object} source
+ * @return {Object}
+ */
+const cloneJSON = (source: any): any => {
+  return JSON.parse(JSON.stringify(source));
+};
+
+/**
+ * returns a result of deep merge of two objects
+ *
+ * @param {Object} target
+ * @param {Object} source
+ * @return {Object}
+ */
+const merge = (
+  target: Record<string, unknown>,
+  source: Record<string, unknown>
+) => {
+  target = cloneJSON(target);
+
+  for (const key in source) {
+    if (source.hasOwnProperty(key)) {
+      const sourceKeyValue = source[key];
+      const targetKeyValue = target[key];
+
+      if (isObject(sourceKeyValue) && isObject(targetKeyValue)) {
+        target[key] = merge(targetKeyValue, sourceKeyValue);
+      } else {
+        target[key] = sourceKeyValue;
+      }
+    }
+  }
+  return target;
+};
+
+/**
+ * get object by reference. works only with local references that points on
+ * definitions object
+ *
+ * @param {String} path
+ * @param {Object} definitions
+ * @return {Object}
+ */
+const getLocalRef = function (
+  inputPath: string,
+  definitions: Record<string, unknown>
+) {
+  const path = inputPath.replace(/^#\/definitions\//, "").split("/");
+
+  const find = function (path: string[], root: any): any {
+    const key = path.shift();
+    if (!key) return {};
+
+    if (!root[key]) {
+      return {};
+    } else if (!path.length) {
+      return root[key];
+    } else {
+      return find(path, root[key]);
+    }
+  };
+
+  const result = find(path, definitions);
+
+  if (!isObject(result)) {
+    return result;
+  }
+  return cloneJSON(result);
+};
+
+/**
+ * merge list of objects from allOf properties
+ * if some of objects contains $ref field extracts this reference and merge it
+ *
+ * @param {Array} allOfList
+ * @param {Object} definitions
+ * @return {Object}
+ */
+const mergeAllOf = function (allOfList: any[], definitions: any) {
+  const length = allOfList.length;
+  let index = -1,
+    result = {};
+
+  while (++index < length) {
+    let item = allOfList[index];
+
+    item =
+      typeof item.$ref !== "undefined"
+        ? getLocalRef(item.$ref, definitions)
+        : item;
+
+    result = merge(result, item);
+  }
+
+  return result;
+};
+
+/**
+ * returns a object that built with default values from json schema
+ *
+ * @param {Object} schema
+ * @param {Object} definitions
+ * @return {Object}
+ */
+const defaults = (schema: any, definitions: Record<string, any>): unknown => {
+  if (typeof schema["default"] !== "undefined") {
+    return schema["default"];
+  } else if (typeof schema.allOf !== "undefined") {
+    const mergedItem = mergeAllOf(schema.allOf, definitions);
+    return defaults(mergedItem, definitions);
+  } else if (typeof schema.$ref !== "undefined") {
+    const reference = getLocalRef(schema.$ref, definitions);
+    return defaults(reference, definitions);
+  } else if (schema.type === "object") {
+    if (!schema.properties) {
+      return {};
+    }
+
+    for (const key in schema.properties) {
+      if (schema.properties.hasOwnProperty(key)) {
+        schema.properties[key] = defaults(schema.properties[key], definitions);
+
+        if (typeof schema.properties[key] === "undefined") {
+          delete schema.properties[key];
+        }
+      }
+    }
+
+    return schema.properties;
+  } else if (schema.type === "array") {
+    if (!schema.items) {
+      return [];
+    }
+
+    // minimum item count
+    const ct = schema.minItems || 0;
+
+    // tuple-typed arrays
+    if (schema.items.constructor === Array) {
+      const values = schema.items.map((item: unknown) =>
+        defaults(item, definitions)
+      );
+
+      // remove undefined items at the end (unless required by minItems)
+      for (let i = values.length - 1; i >= 0; i--) {
+        if (typeof values[i] !== "undefined") {
+          break;
+        }
+        if (i + 1 > ct) {
+          values.pop();
+        }
+      }
+
+      // if all values are undefined -> return undefined even
+      // if minItems is set
+      if (values.every((item: unknown) => typeof item === "undefined")) {
+        return undefined;
+      }
+
+      return values;
+    }
+
+    // object-typed arrays
+    const value = defaults(schema.items, definitions);
+
+    if (typeof value === "undefined") {
+      return [];
+    } else {
+      const values = [];
+      for (let i = 0; i < Math.max(1, ct); i++) {
+        values.push(cloneJSON(value));
+      }
+      return values;
+    }
+  }
+};
+
+/**
+ * main function
+ *
+ * @param {Object} schema
+ * @param {Object|undefined} definitions
+ * @return {Object}
+ */
+export default function (
+  schema: any,
+  definitions?: Record<string, unknown> | undefined
+) {
+  if (typeof definitions === "undefined") {
+    definitions = (schema.definitions as Record<string, unknown>) || {};
+  } else if (isObject(schema.definitions)) {
+    definitions = merge(definitions, schema.definitions);
+  }
+
+  return defaults(cloneJSON(schema), definitions);
+}

langserve/playground/src/utils/schema.ts

@@ -0,0 +1,27 @@
+import { JsonSchema } from "@jsonforms/core";
+
+type JsonSchemaExtra = JsonSchema & {
+  extra: {
+    widget: {
+      type: string;
+    };
+  };
+};
+
+export function isJsonSchemaExtra(x: JsonSchema): x is JsonSchemaExtra {
+  if (!("extra" in x && typeof x.extra === "object" && x.extra != null)) {
+    return false;
+  }
+
+  if (
+    !(
+      "widget" in x.extra &&
+      typeof x.extra.widget === "object" &&
+      x.extra.widget != null
+    )
+  ) {
+    return false;
+  }
+
+  return true;
+}

langserve/playground/src/utils/str.ts

@@ -0,0 +1,5 @@
+export function str(o: unknown): React.ReactNode {
+  return typeof o === "object"
+    ? JSON.stringify(o, null, 2)
+    : (o as React.ReactNode);
+}

langserve/playground/src/utils/url.ts

@@ -1,8 +1,5 @@
 export function resolveApiUrl(path: string) {
-  if (import.meta.env.DEV) {
-    return new URL(path, "http://127.0.0.1:8000");
-  }
-
-  const prefix = window.location.pathname.split("/playground")[0];
+  let prefix = window.location.pathname.split("/playground")[0];
+  if (prefix.endsWith("/")) prefix = prefix.slice(0, -1);
   return new URL(prefix + path, window.location.origin);
 }

langserve/playground/vite.config.ts

@@ -6,4 +6,13 @@ import svgr from "vite-plugin-svgr";
 export default defineConfig({
   base: "/____LANGSERVE_BASE_URL/",
   plugins: [svgr(), react()],
+  server: {
+    proxy: {
+      "^/____LANGSERVE_BASE_URL.*/(config_schema|input_schema|stream_log)": {
+        target: "http://127.0.0.1:8000",
+        changeOrigin: true,
+        rewrite: (path) => path.replace("/____LANGSERVE_BASE_URL", ""),
+      },
+    },
+  },
 });

langserve/playground/yarn.lock

@@ -1224,13 +1224,6 @@ arg@^5.0.2:
   resolved "https://registry.yarnpkg.com/arg/-/arg-5.0.2.tgz#c81433cc427c92c4dcf4865142dbca6f15acd59c"
   integrity sha512-PYjyFOLKQ9y57JvQ6QLo8dAgNqswh8M1RMJYdQduT6xbWSgK36P/Z/v+p888pM69jMMfS8Xd8F6I1kQ/I9HUGg==
 
-argparse@^1.0.9:
-  version "1.0.10"
-  resolved "https://registry.yarnpkg.com/argparse/-/argparse-1.0.10.tgz#bcd6791ea5ae09725e17e5ad988134cd40b3d911"
-  integrity sha512-o5Roy6tNG4SL/FOkCAN6RzjiakZS25RLYFrcMttJqbdd8BWrnA+fGz57iN5Pb06pvBGvl5gQ0B48dJlslXvoTg==
-  dependencies:
-    sprintf-js "~1.0.2"
-
 argparse@^2.0.1:
   version "2.0.1"
   resolved "https://registry.yarnpkg.com/argparse/-/argparse-2.0.1.tgz#246f50f3ca78a3240f6c997e8a9bd1eac49e4b38"
@@ -1985,13 +1978,6 @@ json-parse-even-better-errors@^2.3.0:
   resolved "https://registry.yarnpkg.com/json-parse-even-better-errors/-/json-parse-even-better-errors-2.3.1.tgz#7c47805a94319928e05777405dc12e1f7a4ee02d"
   integrity sha512-xyFwyhro/JEof6Ghe2iz2NcXoj2sloNsWr/XsERDK/oiPCfaNhl5ONfp+jQdAZRQQ0IJWNzH9zIZF7li91kh2w==
 
-json-schema-defaults@^0.4.0:
-  version "0.4.0"
-  resolved "https://registry.yarnpkg.com/json-schema-defaults/-/json-schema-defaults-0.4.0.tgz#b63ee7e7aa83f29b54cb31d31ecddeb056c3306c"
-  integrity sha512-UsUrkDVNvHTneyeQOYHH9ZHb3+6OjwYfJ831SdO0yjtXtYZ7Jh8BKWsuJYUQW7qckP5JhHawsg4GI6A5fMaR/Q==
-  dependencies:
-    argparse "^1.0.9"
-
 json-schema-traverse@^0.4.1:
   version "0.4.1"
   resolved "https://registry.yarnpkg.com/json-schema-traverse/-/json-schema-traverse-0.4.1.tgz#69f6a87d9513ab8bb8fe63bdb0979c448e684660"
@@ -2522,11 +2508,6 @@ source-map@^0.5.7:
   resolved "https://registry.yarnpkg.com/source-map/-/source-map-0.5.7.tgz#8a039d2d1021d22d1ea14c80d8ea468ba2ef3fcc"
   integrity sha512-LbrmJOMUSdEVxIKvdcJzQC+nQhe8FUZQTXQy6+I75skNgn3OoQ0DZA8YnFa7gp8tqtL3KPf1kmo0R5DoApeSGQ==
 
-sprintf-js@~1.0.2:
-  version "1.0.3"
-  resolved "https://registry.yarnpkg.com/sprintf-js/-/sprintf-js-1.0.3.tgz#04e6926f662895354f3dd015203633b857297e2c"
-  integrity sha512-D9cPgkvLlV3t3IzL0D0YLvGA9Ahk4PcvVwUbN0dSGr1aP0Nrt4AEnTUbuGvquEC0mA64Gqt1fzirlRs5ibXx8g==
-
 strip-ansi@^6.0.1:
   version "6.0.1"
   resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9"
@@ -2699,6 +2680,11 @@ use-callback-ref@^1.3.0:
   dependencies:
     tslib "^2.0.0"
 
+use-debounce@^9.0.4:
+  version "9.0.4"
+  resolved "https://registry.yarnpkg.com/use-debounce/-/use-debounce-9.0.4.tgz#51d25d856fbdfeb537553972ce3943b897f1ac85"
+  integrity sha512-6X8H/mikbrt0XE8e+JXRtZ8yYVvKkdYRfmIhWZYsP8rcNs9hk3APV8Ua2mFkKRLcJKVdnX2/Vwrmg2GWKUQEaQ==
+
 use-sidecar@^1.1.2:
   version "1.1.2"
   resolved "https://registry.yarnpkg.com/use-sidecar/-/use-sidecar-1.1.2.tgz#2f43126ba2d7d7e117aa5855e5d8f0276dfe73c2"

langserve/schema.py

@@ -0,0 +1,23 @@
+try:
+    from pydantic.v1 import BaseModel
+except ImportError:
+    from pydantic import BaseModel
+
+
+class CustomUserType(BaseModel):
+    """Inherit from this class to create a custom user type.
+
+    Use a custom user type if you want the data to de-serialize
+    into a pydantic model rather than the equivalent dict representation.
+
+    In general, make sure to add a `type` attribute to your class
+    to help pydantic to discriminate unions.
+
+    https://docs.pydantic.dev/1.10/usage/types/#discriminated-unions-aka-tagged-unions
+
+    Limitations:
+        At the moment, this type only works SERVER side and is used
+        to specify desired DECODING behavior. If inheriting from this type
+        the server will keep the decoded type as a pydantic model instead
+        of converting it into a dict.
+    """

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,12 @@ 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__
+from langserve.schema import CustomUserType
 
 try:
     from pydantic.v1 import BaseModel, create_model
@@ -36,7 +37,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 +46,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
@@ -93,7 +95,9 @@ def _unpack_input(validated_model: BaseModel) -> Any:
     else:
         model = validated_model
 
-    if isinstance(model, BaseModel) and not isinstance(model, Serializable):
+    if isinstance(model, BaseModel) and not isinstance(
+        model, (Serializable, CustomUserType)
+    ):
         # If the model is a pydantic model, but not a Serializable, then
         # it was created by the server as part of validation and isn't expected
         # to be accepted by the runnables as input as a pydantic model,
@@ -199,7 +203,72 @@ 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 _setup_global_app_handlers(app: Union[FastAPI, APIRouter]) -> None:
+    @app.on_event("startup")
+    async def startup_event():
+        # ruff: noqa: E501
+        LANGSERVE = """
+ __          ___      .__   __.   _______      _______. _______ .______     ____    ____  _______ 
+|  |        /   \     |  \ |  |  /  _____|    /       ||   ____||   _  \    \   \  /   / |   ____|
+|  |       /  ^  \    |   \|  | |  |  __     |   (----`|  |__   |  |_)  |    \   \/   /  |  |__   
+|  |      /  /_\  \   |  . `  | |  | |_ |     \   \    |   __|  |      /      \      /   |   __|  
+|  `----./  _____  \  |  |\   | |  |__| | .----)   |   |  |____ |  |\  \----.  \    /    |  |____ 
+|_______/__/     \__\ |__| \__|  \______| |_______/    |_______|| _| `._____|   \__/     |_______|
+"""
+
+        def green(text):
+            return "\x1b[1;32;40m" + text + "\x1b[0m"
+
+        paths = _APP_TO_PATHS[app]
+        print(LANGSERVE)
+        for path in paths:
+            print(
+                f'{green("LANGSERVE:")} Playground for chain "{path or "/"}" is live at:'
+            )
+            print(f'{green("LANGSERVE:")}  │')
+            print(f'{green("LANGSERVE:")}  └──> {path}/playground')
+            print(f'{green("LANGSERVE:")}')
+        print(f'{green("LANGSERVE:")} See all available routes at {app.docs_url}')
+        print()
+
+
+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:
+        _setup_global_app_handlers(app)
+        _APP_TO_PATHS[app] = {path}
 
 
 # PUBLIC API
@@ -213,6 +282,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 +304,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 +327,12 @@ def add_routes(
             "Use `pip install sse_starlette` to install."
         )
 
+    if isinstance(app, FastAPI):  # type: ignore
+        # Cannot do this checking logic for a router since
+        # API routers are not hashable
+        _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 +342,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 +380,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 +415,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 +449,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")
@@ -382,6 +530,16 @@ def add_routes(
                 }
             }
 
+        * error - for signaling an error in the stream, also ends the stream.
+
+        {
+            "event": "error",
+            "data": {
+                "status_code": 500,
+                "message": "Internal Server Error"
+            }
+        }
+
         * end - for signaling the end of the stream.
 
             This helps the client to know when to stop listening for events and
@@ -402,12 +560,27 @@ def add_routes(
 
         async def _stream() -> AsyncIterator[dict]:
             """Stream the output of the runnable."""
-            async for chunk in runnable.astream(
-                input_,
-                config=config,
-            ):
-                yield {"data": simple_dumps(chunk), "event": "data"}
-            yield {"event": "end"}
+            try:
+                async for chunk in runnable.astream(
+                    input_,
+                    config=config,
+                ):
+                    yield {
+                        "data": well_known_lc_serializer.dumps(chunk),
+                        "event": "data",
+                    }
+                yield {"event": "end"}
+            except BaseException:
+                yield {
+                    "event": "error",
+                    # Do not expose the error message to the client since
+                    # the message may contain sensitive information.
+                    # We'll add client side errors for validation as well.
+                    "data": json.dumps(
+                        {"status_code": 500, "message": "Internal Server Error"}
+                    ),
+                }
+                raise
 
         return EventSourceResponse(_stream())
 
@@ -441,6 +614,16 @@ def add_routes(
                 }
             }
 
+        * error - for signaling an error in the stream, also ends the stream.
+
+        {
+            "event": "error",
+            "data": {
+                "status_code": 500,
+                "message": "Internal Server Error"
+            }
+        }
+
         * end - for signaling the end of the stream.
 
             This helps the client to know when to stop listening for events and
@@ -464,31 +647,43 @@ def add_routes(
 
         async def _stream_log() -> AsyncIterator[dict]:
             """Stream the output of the runnable."""
-            async for chunk in runnable.astream_log(
-                input_,
-                config=config,
-                diff=True,
-                include_names=stream_log_request.include_names,
-                include_types=stream_log_request.include_types,
-                include_tags=stream_log_request.include_tags,
-                exclude_names=stream_log_request.exclude_names,
-                exclude_types=stream_log_request.exclude_types,
-                exclude_tags=stream_log_request.exclude_tags,
-            ):
-                if not isinstance(chunk, RunLogPatch):
-                    raise AssertionError(
-                        f"Expected a RunLog instance got {type(chunk)}"
-                    )
-                data = {
-                    "ops": chunk.ops,
-                }
+            try:
+                async for chunk in runnable.astream_log(
+                    input_,
+                    config=config,
+                    diff=True,
+                    include_names=stream_log_request.include_names,
+                    include_types=stream_log_request.include_types,
+                    include_tags=stream_log_request.include_tags,
+                    exclude_names=stream_log_request.exclude_names,
+                    exclude_types=stream_log_request.exclude_types,
+                    exclude_tags=stream_log_request.exclude_tags,
+                ):
+                    if not isinstance(chunk, RunLogPatch):
+                        raise AssertionError(
+                            f"Expected a RunLog instance got {type(chunk)}"
+                        )
+                    data = {
+                        "ops": chunk.ops,
+                    }
 
-                # Temporary adapter
+                    # Temporary adapter
+                    yield {
+                        "data": well_known_lc_serializer.dumps(data),
+                        "event": "data",
+                    }
+                yield {"event": "end"}
+            except BaseException:
                 yield {
-                    "data": simple_dumps(data),
-                    "event": "data",
+                    "event": "error",
+                    # Do not expose the error message to the client since
+                    # the message may contain sensitive information.
+                    # We'll add client side errors for validation as well.
+                    "data": json.dumps(
+                        {"status_code": 500, "message": "Internal Server Error"}
+                    ),
                 }
-            yield {"event": "end"}
+                raise
 
         return EventSourceResponse(_stream_log())
 
@@ -496,33 +691,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(
+        return runnable.get_output_schema(
             _unpack_config(config_hash, keys=config_keys, model=ConfigPayload)
-        ).output_schema.schema()
+        ).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}",
@@ -532,15 +718,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

@@ -827,20 +827,20 @@ tests = ["asttokens (>=2.1.0)", "coverage", "coverage-enable-subprocess", "ipyth
 
 [[package]]
 name = "fastapi"
-version = "0.103.2"
+version = "0.104.0"
 description = "FastAPI framework, high performance, easy to learn, fast to code, ready for production"
 optional = false
-python-versions = ">=3.7"
+python-versions = ">=3.8"
 files = [
-    {file = "fastapi-0.103.2-py3-none-any.whl", hash = "sha256:3270de872f0fe9ec809d4bd3d4d890c6d5cc7b9611d721d6438f9dacc8c4ef2e"},
-    {file = "fastapi-0.103.2.tar.gz", hash = "sha256:75a11f6bfb8fc4d2bec0bd710c2d5f2829659c0e8c0afd5560fdda6ce25ec653"},
+    {file = "fastapi-0.104.0-py3-none-any.whl", hash = "sha256:456482c1178fb7beb2814b88e1885bc49f9a81f079665016feffe3e1c6a7663e"},
+    {file = "fastapi-0.104.0.tar.gz", hash = "sha256:9c44de45693ae037b0c6914727a29c49a40668432b67c859a87851fc6a7b74c6"},
 ]
 
 [package.dependencies]
 anyio = ">=3.7.1,<4.0.0"
 pydantic = ">=1.7.4,<1.8 || >1.8,<1.8.1 || >1.8.1,<2.0.0 || >2.0.0,<2.0.1 || >2.0.1,<2.1.0 || >2.1.0,<3.0.0"
 starlette = ">=0.27.0,<0.28.0"
-typing-extensions = ">=4.5.0"
+typing-extensions = ">=4.8.0"
 
 [package.extras]
 all = ["email-validator (>=2.0.0)", "httpx (>=0.23.0)", "itsdangerous (>=1.1.0)", "jinja2 (>=2.11.2)", "orjson (>=3.2.1)", "pydantic-extra-types (>=2.0.0)", "pydantic-settings (>=2.0.0)", "python-multipart (>=0.0.5)", "pyyaml (>=5.3.1)", "ujson (>=4.0.1,!=4.0.2,!=4.1.0,!=4.2.0,!=4.3.0,!=5.0.0,!=5.1.0)", "uvicorn[standard] (>=0.12.0)"]
@@ -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.316"
+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.316-py3-none-any.whl", hash = "sha256:997d2ecdaf481517c90dcea86f93b98327546c357ac8f581f06d7d782a2eda2d"},
-    {file = "langchain-0.0.316.tar.gz", hash = "sha256:732c958d41e22a5bb55f8c0cb70914793f083d175739d476e5cbbc7715703153"},
+    {file = "langchain-0.0.322-py3-none-any.whl", hash = "sha256:f065d19eff2702cd941fd1f6adca6baf2a1181387f04d3429a0ec7a197e96155"},
+    {file = "langchain-0.0.322.tar.gz", hash = "sha256:8415327a964bff9d643202697aca1dcbcfd2d89fb0b49f799bdea37d01fa7f51"},
 ]
 
 [package.dependencies]
@@ -1685,7 +1686,7 @@ cli = ["typer (>=0.9.0,<0.10.0)"]
 cohere = ["cohere (>=4,<5)"]
 docarray = ["docarray[hnswlib] (>=0.32.0,<0.33.0)"]
 embeddings = ["sentence-transformers (>=2,<3)"]
-extended-testing = ["aiosqlite (>=0.19.0,<0.20.0)", "amazon-textract-caller (<2)", "anthropic (>=0.3.11,<0.4.0)", "arxiv (>=1.4,<2.0)", "assemblyai (>=0.17.0,<0.18.0)", "atlassian-python-api (>=3.36.0,<4.0.0)", "beautifulsoup4 (>=4,<5)", "bibtexparser (>=1.4.0,<2.0.0)", "cassio (>=0.1.0,<0.2.0)", "chardet (>=5.1.0,<6.0.0)", "dashvector (>=1.0.1,<2.0.0)", "esprima (>=4.0.1,<5.0.0)", "faiss-cpu (>=1,<2)", "feedparser (>=6.0.10,<7.0.0)", "geopandas (>=0.13.1,<0.14.0)", "gitpython (>=3.1.32,<4.0.0)", "gql (>=3.4.1,<4.0.0)", "html2text (>=2020.1.16,<2021.0.0)", "jinja2 (>=3,<4)", "jq (>=1.4.1,<2.0.0)", "lxml (>=4.9.2,<5.0.0)", "markdownify (>=0.11.6,<0.12.0)", "motor (>=3.3.1,<4.0.0)", "mwparserfromhell (>=0.6.4,<0.7.0)", "mwxml (>=0.3.3,<0.4.0)", "newspaper3k (>=0.2.8,<0.3.0)", "numexpr (>=2.8.6,<3.0.0)", "openai (>=0,<1)", "openapi-schema-pydantic (>=1.2,<2.0)", "pandas (>=2.0.1,<3.0.0)", "pdfminer-six (>=20221105,<20221106)", "pgvector (>=0.1.6,<0.2.0)", "psychicapi (>=0.8.0,<0.9.0)", "py-trello (>=0.19.0,<0.20.0)", "pymupdf (>=1.22.3,<2.0.0)", "pypdf (>=3.4.0,<4.0.0)", "pypdfium2 (>=4.10.0,<5.0.0)", "pyspark (>=3.4.0,<4.0.0)", "rank-bm25 (>=0.2.2,<0.3.0)", "rapidfuzz (>=3.1.1,<4.0.0)", "rapidocr-onnxruntime (>=1.3.2,<2.0.0)", "requests-toolbelt (>=1.0.0,<2.0.0)", "rspace_client (>=2.5.0,<3.0.0)", "scikit-learn (>=1.2.2,<2.0.0)", "sqlite-vss (>=0.1.2,<0.2.0)", "streamlit (>=1.18.0,<2.0.0)", "sympy (>=1.12,<2.0)", "telethon (>=1.28.5,<2.0.0)", "timescale-vector (>=0.0.1,<0.0.2)", "tqdm (>=4.48.0)", "upstash-redis (>=0.15.0,<0.16.0)", "xata (>=1.0.0a7,<2.0.0)", "xmltodict (>=0.13.0,<0.14.0)"]
+extended-testing = ["aiosqlite (>=0.19.0,<0.20.0)", "amazon-textract-caller (<2)", "anthropic (>=0.3.11,<0.4.0)", "arxiv (>=1.4,<2.0)", "assemblyai (>=0.17.0,<0.18.0)", "atlassian-python-api (>=3.36.0,<4.0.0)", "beautifulsoup4 (>=4,<5)", "bibtexparser (>=1.4.0,<2.0.0)", "cassio (>=0.1.0,<0.2.0)", "chardet (>=5.1.0,<6.0.0)", "dashvector (>=1.0.1,<2.0.0)", "esprima (>=4.0.1,<5.0.0)", "faiss-cpu (>=1,<2)", "feedparser (>=6.0.10,<7.0.0)", "geopandas (>=0.13.1,<0.14.0)", "gitpython (>=3.1.32,<4.0.0)", "gql (>=3.4.1,<4.0.0)", "html2text (>=2020.1.16,<2021.0.0)", "jinja2 (>=3,<4)", "jq (>=1.4.1,<2.0.0)", "lxml (>=4.9.2,<5.0.0)", "markdownify (>=0.11.6,<0.12.0)", "motor (>=3.3.1,<4.0.0)", "mwparserfromhell (>=0.6.4,<0.7.0)", "mwxml (>=0.3.3,<0.4.0)", "newspaper3k (>=0.2.8,<0.3.0)", "numexpr (>=2.8.6,<3.0.0)", "openai (>=0,<1)", "openapi-pydantic (>=0.3.2,<0.4.0)", "pandas (>=2.0.1,<3.0.0)", "pdfminer-six (>=20221105,<20221106)", "pgvector (>=0.1.6,<0.2.0)", "psychicapi (>=0.8.0,<0.9.0)", "py-trello (>=0.19.0,<0.20.0)", "pymupdf (>=1.22.3,<2.0.0)", "pypdf (>=3.4.0,<4.0.0)", "pypdfium2 (>=4.10.0,<5.0.0)", "pyspark (>=3.4.0,<4.0.0)", "rank-bm25 (>=0.2.2,<0.3.0)", "rapidfuzz (>=3.1.1,<4.0.0)", "rapidocr-onnxruntime (>=1.3.2,<2.0.0)", "requests-toolbelt (>=1.0.0,<2.0.0)", "rspace_client (>=2.5.0,<3.0.0)", "scikit-learn (>=1.2.2,<2.0.0)", "sqlite-vss (>=0.1.2,<0.2.0)", "streamlit (>=1.18.0,<2.0.0)", "sympy (>=1.12,<2.0)", "telethon (>=1.28.5,<2.0.0)", "timescale-vector (>=0.0.1,<0.0.2)", "tqdm (>=4.48.0)", "upstash-redis (>=0.15.0,<0.16.0)", "xata (>=1.0.0a7,<2.0.0)", "xmltodict (>=0.13.0,<0.14.0)"]
 javascript = ["esprima (>=4.0.1,<5.0.0)"]
 llms = ["clarifai (>=9.1.0)", "cohere (>=4,<5)", "huggingface_hub (>=0,<1)", "manifest-ml (>=0.0.1,<0.0.2)", "nlpcloud (>=1,<2)", "openai (>=0,<1)", "openlm (>=0.0.5,<0.0.6)", "torch (>=1,<3)", "transformers (>=4,<5)"]
 openai = ["openai (>=0,<1)", "tiktoken (>=0.3.2,<0.6.0)"]
@@ -1694,13 +1695,13 @@ text-helpers = ["chardet (>=5.1.0,<6.0.0)"]
 
 [[package]]
 name = "langsmith"
-version = "0.0.44"
+version = "0.0.47"
 description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
 optional = false
 python-versions = ">=3.8.1,<4.0"
 files = [
-    {file = "langsmith-0.0.44-py3-none-any.whl", hash = "sha256:5e7e5b45360ce89a2d5d6066a3b9fdd31b1f874a0cf19b1666c9792fecef0a1b"},
-    {file = "langsmith-0.0.44.tar.gz", hash = "sha256:74a262ba23a958ca1a4863d5386c151be462e40ccfcb8b39d0a5d8c9eaa40164"},
+    {file = "langsmith-0.0.47-py3-none-any.whl", hash = "sha256:90279d4888e4513f83703becf38a6d92b7c7d696b120872fa1fb28c5c67bfd91"},
+    {file = "langsmith-0.0.47.tar.gz", hash = "sha256:04b8373cbfcf7b9c28ba53174206095c50dc09ae1131a99b501f92776c8ad0c8"},
 ]
 
 [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]
@@ -2496,13 +2507,13 @@ testing = ["fields", "hunter", "process-tests", "pytest-xdist", "six", "virtuale
 
 [[package]]
 name = "pytest-mock"
-version = "3.11.1"
+version = "3.12.0"
 description = "Thin-wrapper around the mock package for easier use with pytest"
 optional = false
-python-versions = ">=3.7"
+python-versions = ">=3.8"
 files = [
-    {file = "pytest-mock-3.11.1.tar.gz", hash = "sha256:7f6b125602ac6d743e523ae0bfa71e1a697a2f5534064528c6ff84c2f7c2fc7f"},
-    {file = "pytest_mock-3.11.1-py3-none-any.whl", hash = "sha256:21c279fff83d70763b05f8874cc9cfb3fcacd6d354247a976f9529d19f9acf39"},
+    {file = "pytest-mock-3.12.0.tar.gz", hash = "sha256:31a40f038c22cad32287bb43932054451ff5583ff094bca6f675df2f8bc1a6e9"},
+    {file = "pytest_mock-3.12.0-py3-none-any.whl", hash = "sha256:0972719a7263072da3a21c7f4773069bcc7486027d7e8e1f81d98a47e701bc4f"},
 ]
 
 [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"},
@@ -3328,13 +3349,13 @@ dev = ["flake8", "flake8-annotations", "flake8-bandit", "flake8-bugbear", "flake
 
 [[package]]
 name = "urllib3"
-version = "2.0.6"
+version = "2.0.7"
 description = "HTTP library with thread-safe connection pooling, file post, and more."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "urllib3-2.0.6-py3-none-any.whl", hash = "sha256:7a7c7003b000adf9e7ca2a377c9688bbc54ed41b985789ed576570342a375cd2"},
-    {file = "urllib3-2.0.6.tar.gz", hash = "sha256:b19e1a85d206b56d7df1d5e683df4a7725252a964e3993648dd0fb5a1c157564"},
+    {file = "urllib3-2.0.7-py3-none-any.whl", hash = "sha256:fdb6d215c776278489906c2f8916e6e7d4f5a9b602ccbcfdf7f016fc8da0596e"},
+    {file = "urllib3-2.0.7.tar.gz", hash = "sha256:c97dfde1f7bd43a71c8d2a58e369e9b2bf692d1334ea9f9cae55add7d0dd0f84"},
 ]
 
 [package.extras]
@@ -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.13"
+version = "0.0.20"
 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

@@ -2,38 +2,41 @@
 import asyncio
 import json
 from asyncio import AbstractEventLoop
-from contextlib import asynccontextmanager
-from typing import Any, Dict, List, Optional, Union
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 import httpx
 import pytest
 import pytest_asyncio
-from fastapi import FastAPI
+from fastapi import APIRouter, FastAPI
 from fastapi.testclient import TestClient
 from httpx import AsyncClient
 from langchain.callbacks.tracers.log_stream import RunLogPatch
 from langchain.prompts import PromptTemplate
 from langchain.prompts.base import StringPromptValue
 from langchain.schema.messages import HumanMessage, SystemMessage
-from langchain.schema.runnable import RunnableConfig, RunnablePassthrough
+from langchain.schema.runnable import Runnable, RunnableConfig, RunnablePassthrough
 from langchain.schema.runnable.base import RunnableLambda
-from langchain.schema.runnable.utils import ConfigurableField
+from langchain.schema.runnable.utils import ConfigurableField, Input, Output
 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.schema import CustomUserType
 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")
@@ -62,7 +65,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
@@ -103,14 +108,19 @@ def client(app: FastAPI) -> RemoteRunnable:
 
 @asynccontextmanager
 async def get_async_client(
-    server: FastAPI, path: Optional[str] = None
+    server: FastAPI, *, path: Optional[str] = None, raise_app_exceptions: bool = True
 ) -> RemoteRunnable:
     """Get an async client."""
     url = "http://localhost:9999"
     if path:
         url += path
     remote_runnable_client = RemoteRunnable(url=url)
-    async_client = AsyncClient(app=server, base_url=url)
+
+    transport = httpx.ASGITransport(
+        app=server,
+        raise_app_exceptions=raise_app_exceptions,
+    )
+    async_client = AsyncClient(app=server, base_url=url, transport=transport)
     remote_runnable_client.async_client = async_client
     try:
         yield remote_runnable_client
@@ -118,6 +128,25 @@ async def get_async_client(
         await async_client.aclose()
 
 
+@contextmanager
+def get_sync_client(
+    server: FastAPI, *, path: Optional[str] = None, raise_server_exceptions: bool = True
+) -> RemoteRunnable:
+    """Get an async client."""
+    url = "http://localhost:9999"
+    if path:
+        url += path
+    remote_runnable_client = RemoteRunnable(url=url)
+    sync_client = TestClient(
+        app=server, base_url=url, raise_server_exceptions=raise_server_exceptions
+    )
+    remote_runnable_client.sync_client = sync_client
+    try:
+        yield remote_runnable_client
+    finally:
+        sync_client.close()
+
+
 @pytest_asyncio.fixture()
 async def async_client(app: FastAPI) -> RemoteRunnable:
     """Create a FastAPI app that exposes the Runnable as an API."""
@@ -131,19 +160,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"
@@ -153,11 +187,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."""
@@ -165,13 +210,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})
@@ -190,8 +238,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
@@ -200,9 +249,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(
@@ -223,6 +272,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."""
@@ -232,15 +292,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:
@@ -251,6 +362,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
@@ -522,7 +672,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
@@ -530,18 +680,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"]
@@ -553,10 +703,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):
@@ -569,15 +721,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:
@@ -590,10 +744,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:
@@ -749,6 +905,128 @@ def test_rename_pydantic_model() -> None:
     assert Model.__name__ == "Bar"
 
 
+@pytest.mark.asyncio
+async def test_input_config_output_schemas(event_loop: AbstractEventLoop) -> None:
+    """Test schemas returned for different configurations."""
+    # TODO(Fix me): need to fix handling of global state -- we get problems
+    # gives inconsistent results when running multiple tests / results
+    # depending on ordering
+    server._SEEN_NAMES = set()
+    server._MODEL_REGISTRY = {}
+
+    async def add_one(x: int) -> int:
+        """Add one to simulate a valid function"""
+        return x + 1
+
+    async def add_two(y: int) -> int:
+        """Add one to simulate a valid function"""
+        return y + 2
+
+    app = FastAPI()
+
+    add_routes(app, RunnableLambda(add_one), path="/add_one")
+    # Custom input type
+    add_routes(
+        app,
+        RunnableLambda(add_two),
+        path="/add_two_custom",
+        input_type=float,
+        output_type=float,
+        config_keys=["tags", "configurable"],
+    )
+    add_routes(app, PromptTemplate.from_template("{question}"), path="/prompt_1")
+
+    template = PromptTemplate.from_template("say {name}").configurable_fields(
+        template=ConfigurableField(
+            id="template",
+            name="Template",
+            description="The template to use for the prompt",
+        )
+    )
+    add_routes(app, template, path="/prompt_2", config_keys=["tags", "configurable"])
+
+    async with AsyncClient(app=app, base_url="http://localhost:9999") as async_client:
+        # input schema
+        response = await async_client.get("/add_one/input_schema")
+        assert response.json() == {"title": "RunnableLambdaInput", "type": "integer"}
+
+        response = await async_client.get("/add_two_custom/input_schema")
+        assert response.json() == {"title": "RunnableBindingInput", "type": "number"}
+
+        response = await async_client.get("/prompt_1/input_schema")
+        assert response.json() == {
+            "properties": {"question": {"title": "Question", "type": "string"}},
+            "title": "PromptInput",
+            "type": "object",
+        }
+
+        response = await async_client.get("/prompt_2/input_schema")
+        assert response.json() == {
+            "properties": {"name": {"title": "Name", "type": "string"}},
+            "title": "PromptInput",
+            "type": "object",
+        }
+
+        # output schema
+        response = await async_client.get("/add_one/output_schema")
+        assert response.json() == {
+            "title": "RunnableLambdaOutput",
+            "type": "integer",
+        }
+
+        response = await async_client.get("/add_two_custom/output_schema")
+        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
+        response = await async_client.get("/prompt_1/output_schema")
+        assert response.json() != {}  # Long string
+
+        response = await async_client.get("/prompt_2/output_schema")
+        assert response.json() != {}  # Long string
+
+        ## Config schema
+        response = await async_client.get("/add_one/config_schema")
+        assert response.json() == {
+            "properties": {},
+            "title": "RunnableLambdaConfig",
+            "type": "object",
+        }
+
+        response = await async_client.get("/add_two_custom/config_schema")
+        assert response.json() == {
+            "properties": {
+                "tags": {"items": {"type": "string"}, "title": "Tags", "type": "array"}
+            },
+            "title": "RunnableLambdaConfig",
+            "type": "object",
+        }
+
+        response = await async_client.get("/prompt_2/config_schema")
+        assert response.json() == {
+            "definitions": {
+                "Configurable": {
+                    "properties": {
+                        "template": {
+                            "default": "say {name}",
+                            "description": "The template to use for the prompt",
+                            "title": "Template",
+                            "type": "string",
+                        }
+                    },
+                    "title": "Configurable",
+                    "type": "object",
+                }
+            },
+            "properties": {
+                "configurable": {"$ref": "#/definitions/Configurable"},
+                "tags": {"items": {"type": "string"}, "title": "Tags", "type": "array"},
+            },
+            "title": "RunnableConfigurableFieldsConfig",
+            "type": "object",
+        }
+
+
 @pytest.mark.asyncio
 async def test_input_schema_typed_dict() -> None:
     class InputType(TypedDict):
@@ -765,8 +1043,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": {
@@ -782,4 +1059,183 @@ async def test_input_schema_typed_dict() -> None:
                     "type": "object",
                 }
             },
+            "title": "RunnableBindingInput",
         }
+
+
+class ErroringRunnable(Runnable):
+    """A custom runnable for testing errors are raised server side."""
+
+    def invoke(self, input: Input, config: Optional[RunnableConfig] = None) -> Output:
+        """Invoke the runnable."""
+        raise ValueError("Server side error")
+
+    def stream(
+        self,
+        input: Input,
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Optional[Any],
+    ) -> Iterator[Output]:
+        yield 1
+        yield 2
+        raise ValueError("An exception occurred")
+
+    async def astream(
+        self,
+        input: Iterator[Input],
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Optional[Any],
+    ) -> Iterator[Output]:
+        yield 1
+        yield 2
+        raise ValueError("An exception occurred")
+
+
+@pytest.mark.asyncio
+async def test_server_side_error() -> None:
+    """Test server side error handling."""
+
+    app = FastAPI()
+    add_routes(app, ErroringRunnable())
+
+    # 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, 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], 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:
+            async for chunk in runnable.astream({"a": 1}):
+                chunks.append(chunk)
+        except httpx.HTTPStatusError as e:
+            assert chunks == [1, 2]
+            assert e.response.status_code == 500
+            assert e.response.text == "Internal Server Error"
+
+        # # Failing right now, can uncomment or add callbacks
+        # # Test astream_log
+        # chunks = []
+        # try:
+        #     async for chunk in runnable.astream_log({"a": 1}):
+        #         chunks.append(chunk)
+        # except httpx.HTTPStatusError as e:
+        #     assert chunks == []
+        #     assert e.response.status_code == 500
+        #     assert e.response.text == "Internal Server Error"
+
+
+def test_server_side_error_sync() -> None:
+    """Test server side error handling."""
+
+    app = FastAPI()
+    add_routes(app, ErroringRunnable())
+
+    # Invoke request
+    with get_sync_client(app, raise_server_exceptions=False) as runnable:
+        with pytest.raises(httpx.HTTPStatusError) as cm:
+            assert runnable.invoke(1)
+        assert isinstance(cm.value, httpx.HTTPStatusError)
+
+        with pytest.raises(httpx.HTTPStatusError) as cm:
+            assert runnable.batch([1, 2])
+        assert isinstance(cm.value, httpx.HTTPStatusError)
+
+        # Test astream
+        chunks = []
+        try:
+            for chunk in runnable.stream({"a": 1}):
+                chunks.append(chunk)
+        except httpx.HTTPStatusError as e:
+            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")
+    with pytest.raises(ValueError):
+        add_routes(app, RunnableLambda(lambda foo: "hello"), path="/foo")
+    add_routes(app, RunnableLambda(lambda foo: "hello"), path="/baz")
+
+
+@pytest.mark.asyncio
+async def test_custom_user_type() -> None:
+    """Test custom user type."""
+    app = FastAPI()
+
+    class Foo(CustomUserType):
+        bar: int
+
+    def func(foo: Foo) -> int:
+        """Sample function that expects a Foo type which is a pydantic model"""
+        assert isinstance(foo, Foo)
+        return foo.bar
+
+    class Baz(BaseModel):
+        bar: int
+
+    def func2(baz) -> int:
+        """Sample function that expects a Foo type which is a pydantic model"""
+        assert isinstance(baz, dict)
+        return baz["bar"]
+
+    add_routes(app, RunnableLambda(func), path="/foo")
+    add_routes(app, RunnableLambda(func2).with_types(input_type=Baz), path="/baz")
+
+    # Invoke request
+    async with get_async_client(
+        app, path="/foo", raise_app_exceptions=False
+    ) as runnable:
+        assert await runnable.ainvoke({"bar": 1}) == 1
+
+
+@pytest.mark.asyncio
+async def test_using_router() -> None:
+    """Test using a router."""
+    app = FastAPI()
+
+    # Make sure that we can add routers
+    # to an API router
+    router = APIRouter()
+
+    add_routes(
+        router,
+        RunnableLambda(lambda foo: "hello"),
+        path="/chat",
+    )
+
+    app.include_router(router)

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))