From f0ed9830b01711e0c497efd6796beb5ba4af91e1 Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Thu, 16 Nov 2023 10:10:29 -0500
Subject: [PATCH] Add support for API route via prefix (#236)

* Fix serving of playground assets when using an API router with a
custom prefix

---------

Co-authored-by: Andreas Hildebrandt <andreas.hildebrandt@uni-mainz.de>
---
 langserve/server.py                    |  8 +++++++-
 tests/unit_tests/test_server_client.py | 21 +++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/langserve/server.py b/langserve/server.py
index 6dcabf61..c7b075d4 100644
--- a/langserve/server.py
+++ b/langserve/server.py
@@ -1025,11 +1025,17 @@ async def playground(
                 request=request,
                 per_req_config_modifier=per_req_config_modifier,
             )
+
+        if isinstance(app, FastAPI):  # type: ignore
+            base_url = f"{namespace}/playground"
+        else:
+            base_url = f"{app.prefix}{namespace}/playground"
+
         return await serve_playground(
             runnable.with_config(config),
             runnable.with_config(config).input_schema,
             config_keys,
-            f"{namespace}/playground",
+            base_url,
             file_path,
         )
 
diff --git a/tests/unit_tests/test_server_client.py b/tests/unit_tests/test_server_client.py
index 3b943698..b18d9c81 100644
--- a/tests/unit_tests/test_server_client.py
+++ b/tests/unit_tests/test_server_client.py
@@ -269,6 +269,27 @@ def test_serve_playground(app: FastAPI) -> None:
     assert response.status_code == 404
 
 
+@pytest.mark.asyncio
+async def test_serve_playground_with_api_router() -> None:
+    """Test serving playground from an api router with a prefix."""
+    app = FastAPI()
+
+    # Make sure that we can add routers
+    # to an API router
+    router = APIRouter(prefix="/langserve_runnables")
+
+    add_routes(
+        router,
+        RunnableLambda(lambda foo: "hello"),
+        path="/chat",
+    )
+
+    app.include_router(router)
+    async_client = AsyncClient(app=app, base_url="http://localhost:9999")
+    response = await async_client.get("/langserve_runnables/chat/playground/index.html")
+    assert response.status_code == 200
+
+
 @pytest.mark.asyncio
 async def test_server_async(app: FastAPI) -> None:
     """Test the server directly via HTTP requests."""