Merge branch 'main' into jacob/allowlist

langserve/client.py

@@ -133,7 +133,11 @@ class RemoteRunnable(Runnable[Input, Output]):
                            and async httpx clients
         """
         _client_kwargs = client_kwargs or {}
-        self.url = url
+        # need to confirm that the url ends with a trailing slash
+        # otherwise urljoin will remove the last segment, assuming
+        # it's a filepath
+
+        self.url = url if url.endswith("/") else url + "/"
         self.sync_client = httpx.Client(
             base_url=url,
             timeout=timeout,
@@ -454,6 +458,8 @@ class RemoteRunnable(Runnable[Input, Output]):
             async with aconnect_sse(
                 self.async_client, "POST", endpoint, json=data
             ) as event_source:
+                event_source.response.raise_for_status()
+
                 async for sse in event_source.aiter_sse():
                     if sse.event == "data":
                         data = simple_loads(sse.data)

langserve/server.py

@@ -14,6 +14,7 @@ from typing import (
     Dict,
     Literal,
     Mapping,
+    Optional,
     Sequence,
     Type,
     Union,
@@ -208,6 +209,7 @@ def add_routes(
     input_type: Union[Type, Literal["auto"], BaseModel] = "auto",
     output_type: Union[Type, Literal["auto"], BaseModel] = "auto",
     config_keys: Sequence[str] = (),
+    stream_log_name_allowlist: Optional[Sequence[str]] = None,
 ) -> None:
     """Register the routes on the given FastAPI app or APIRouter.
 
@@ -233,7 +235,9 @@ def add_routes(
             Default is "auto" which will use the OutputType of the runnable.
             User is free to provide a custom type annotation.
         config_keys: list of config keys that will be accepted, by default
-                     no config keys are accepted.
+            no config keys are accepted.
+        stream_log_name_allowlist: list of run names that the client can
+            stream as intermediate steps
     """
     try:
         from sse_starlette import EventSourceResponse
@@ -334,7 +338,7 @@ def add_routes(
         request: Request,
         config_hash: str = "",
     ) -> EventSourceResponse:
-        """Invoke the runnable stream the output.
+        """Invoke the runnable stream method.
 
         This endpoint allows to stream the output of the runnable.
 
@@ -392,7 +396,7 @@ def add_routes(
         request: Request,
         config_hash: str = "",
     ) -> EventSourceResponse:
-        """Invoke the runnable stream_log the output.
+        """Invoke the runnable stream_log method.
 
         This endpoint allows to stream the output of the runnable, including
         the output of all intermediate steps.
@@ -424,6 +428,38 @@ def add_routes(
                 "event": "end",
             }
         """
+
+        def _check_allowlist(
+            stream_log_request: Annotated[StreamLogRequest, StreamLogRequest],
+            name_allowlist: Optional[Sequence[str]],
+        ) -> None:
+            if name_allowlist is None:
+                return
+            elif bool(stream_log_request.include_types):
+                raise HTTPException(
+                    403,
+                    "The `include_types` parameter is not permitted due to a run name allowlist set on this endpoint.",
+                )
+            elif bool(stream_log_request.include_tags):
+                raise HTTPException(
+                    403,
+                    "The `include_tags` parameter is not permitted due to a run name allowlist set on this endpoint.",
+                )
+            elif bool(stream_log_request.include_names):
+                if bool(name_allowlist):
+                    allowed_names_string = ", ".join(name_allowlist)
+                    raise HTTPException(
+                        403,
+                        f"Some names in `include_names` are not permitted due to a run name allowlist set on this endpoint.\nAllowed names are: {allowed_names_string}",
+                    )
+                else:
+                    raise HTTPException(
+                        403,
+                        "The `include_names` is not permitted due to a run name allowlist set on this endpoint.",
+                    )
+            return
+
+        _check_allowlist(stream_log_request, stream_log_name_allowlist)
         # Request is first validated using InvokeRequest which takes into account
         # config_keys as well as input_type.
         # After validation, the input is loaded using LangChain's load function.

tests/unit_tests/test_server_client.py

@@ -335,6 +335,109 @@ async def test_astream_log(async_client: RemoteRunnable) -> None:
     ]
 
 
+@pytest.mark.asyncio
+async def test_astream_log_allowlist(event_loop: AbstractEventLoop) -> None:
+    """Test async stream with an allowlist."""
+
+    async def add_one(x: int) -> int:
+        """Add one to simulate a valid function"""
+        return x + 1
+
+    app = FastAPI()
+    add_routes(
+        app,
+        RunnableLambda(add_one),
+        path="/empty_allowlist",
+        input_type=int,
+        stream_log_name_allowlist=["allowed"],
+    )
+    add_routes(
+        app,
+        RunnableLambda(add_one),
+        input_type=int,
+        path="/allowlist",
+        stream_log_name_allowlist=["allowed"],
+    )
+
+    async with get_async_client(app, path="/empty_allowlist") as runnable:
+        with pytest.raises(httpx.HTTPError):
+            async for chunk in runnable.astream_log(1, include_tags=["tag"]):
+                pass
+        with pytest.raises(httpx.HTTPError):
+            async for chunk in runnable.astream_log(1, include_types=["type"]):
+                pass
+        with pytest.raises(httpx.HTTPError):
+            async for chunk in runnable.astream_log(1, include_names=["name"]):
+                pass
+
+        run_log_patches = []
+
+        async for chunk in runnable.astream_log(1):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+        run_log_patches = []
+        async for chunk in runnable.astream_log(1, include_tags=[]):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+        run_log_patches = []
+        async for chunk in runnable.astream_log(1, include_types=[]):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+        run_log_patches = []
+        async for chunk in runnable.astream_log(1, include_names=[]):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+    async with get_async_client(app, path="/allowlist") as runnable:
+        with pytest.raises(httpx.HTTPError):
+            async for chunk in runnable.astream_log(1, include_tags=["tag"]):
+                pass
+        with pytest.raises(httpx.HTTPError):
+            async for chunk in runnable.astream_log(1, include_types=["type"]):
+                pass
+        with pytest.raises(httpx.HTTPError):
+            async for chunk in runnable.astream_log(1, include_names=["name"]):
+                pass
+
+        run_log_patches = []
+
+        async for chunk in runnable.astream_log(1):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+        run_log_patches = []
+        async for chunk in runnable.astream_log(1, include_tags=[]):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+        run_log_patches = []
+        async for chunk in runnable.astream_log(1, include_types=[]):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+        run_log_patches = []
+        async for chunk in runnable.astream_log(1, include_names=[]):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+        run_log_patches = []
+        async for chunk in runnable.astream_log(1, include_names=["allowed"]):
+            run_log_patches.append(chunk)
+
+        assert len(run_log_patches) > 0
+
+
 def test_invoke_as_part_of_sequence(client: RemoteRunnable) -> None:
     """Test as part of sequence."""
     runnable = client | RunnableLambda(func=lambda x: x + 1)