Streaming with the per_req_config_modifier

[rarchitgupta]

I am currently developing an app using langserve that requires a user to be authenticated to load files into the vectorstore and for retrieving documents as well:

Based on the example provided here:
https://github.com/langchain-ai/langserve/blob/main/examples/auth/per_req_config_modifier/server.py

I have come up with this so far which works perfectly for the "invoke" response:

class PerUserQuery(RunnableSerializable):
    """
    A custom runnable that returns a list of documents for the given user.

    The runnable is configurable by the user, and the search results are
    filtered by the user ID.
    """

    user_id: Optional[str]
    vectorstore: VectorStore

    class Config:
        # Allow arbitrary types since VectorStore is an abstract interface
        # and not a pydantic model
        arbitrary_types_allowed = True

    def get_rag_chain(self, query):

        def format_docs(docs):
            return "\n\n".join(doc.page_content for doc in docs)

        retriever = self.vectorstore.similarity_search(
            query, k=10, filter={"user_id": {"$eq": self.user_id}}
        )
        formatted_docs = format_docs(retriever)
        setup_and_retrieval = RunnableParallel(
            {"context": lambda x: formatted_docs, "question": RunnablePassthrough()}
        )
        prompt = hub.pull("rarchit/rag-prompt-bullet")
        llm = get_llm()
        rag_chain = setup_and_retrieval | prompt | llm | StrOutputParser()
        return rag_chain

    def _invoke(
        self, input: str, config: Optional[RunnableConfig] = None, **kwargs: Any
    ) -> List[Document]:
        """
        Invoke the retriever
        """
        rag_chain = self.get_rag_chain(query=input)
        return rag_chain.invoke(str(input))

    def invoke(
        self, input: str, config: Optional[RunnableConfig] = None, **kwargs
    ) -> List[Document]:
        return self._call_with_config(self._invoke, input, config, **kwargs)

However I am unable to define a request modifier for the stream endpoint
I tried:

async def _stream(
    self, input: str, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> AsyncGenerator[str, None]:
    """
    Stream the retriever
    """
    rag_chain = self.get_rag_chain(query=input)
    async for chunk in rag_chain.stream(str(input)):
        yield chunk

async def stream(
    self, input: str, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> StreamingResponse:
    return StreamingResponse(
        self._call_with_config(self._stream, input, config, **kwargs)
    )

But did not have any luck with streaming, it just sends the entire response at once. I would like to obtain the same streaming behaviour as I would if I added the rag_chain directly with:

add_routes(app, rag_chain)

Any help would be appreciated. Stuck here

[ltbd78]

Also doing something similar. Would love to see if anyone has an answer for this. Also curious if it's possible to enable a /playground endpoint for this PerUserVectorstore

[ltbd78]

Yes, so only retriever should exist within your PerUserQuery class (just like the original unmodified code) and the remaining chain prompt | llm | StrOutputParser() should exist outside. There shouldn't be multiple Runnables within your PerUserQuery

Then you do: per_user_query_object | prompt | llm | StrOutputParser()

You also don't need def stream / _stream functions since retriever can't stream.

Then streaming endpoints should work for llm, the playground also works apparently which I was surprised.

[rarchitgupta]

If you don't mind, could you share a snippet of the same? I somewhat get the idea of what you're saying and tried to implement it

[ltbd78]

Unfortunately the code I have is proprietary. But if you follow the steps I listed precisely, it should work. There was a slight mistake I edited.

Essentially it's a merger of two files:
https://github.com/langchain-ai/langserve/blob/main/examples/auth/per_req_config_modifier/server.py
https://github.com/langchain-ai/langserve/blob/main/examples/conversational_retrieval_chain/server.py

[rarchitgupta]

@ltbd78 Thank you so much, yes I dug deeper into what you mentioned earlier and after tinkering around for quite a bit, I got something that works for me

Modified the class to look like this:

class PerUserQuery(RunnableSerializable):
    """
    A custom runnable that returns a list of documents for the given user.

    The runnable is configurable by the user, and the search results are
    filtered by the user ID.
    """

    user_id: Optional[str]
    vectorstore: VectorStore

    class Config:
        # Allow arbitrary types since VectorStore is an abstract interface
        # and not a pydantic model
        arbitrary_types_allowed = True

    def get_rag_chain(self, query):

        retriever = self.vectorstore.similarity_search(
            query, k=10, filter={"user_id": {"$eq": self.user_id}}
        )
        def format_docs(docs):
            return "\n\n".join(doc.page_content for doc in docs)
        
        return {"context": format_docs(retriever), "question": RunnablePassthrough()}

    def invoke(
        self, input: str, config: Optional[RunnableConfig] = None, **kwargs
    ) -> List[Document]:
        retriever = self.get_rag_chain(query=input)
        return retriever

And then set up the final chain like so:

per_user_retriever = PerUserQuery(
    user_id=None,  # Placeholder ID that will be replaced by the per_req_config_modifier
    vectorstore=get_vectorstore(),
).configurable_fields(
    user_id=ConfigurableField(
        id="user_id",
        name="User ID",
        description="The user ID to use for the retriever",
    )
)

prompt = hub.pull("rarchit/rag-prompt-bullet")
llm = get_llm()

chain = per_user_retriever | prompt | llm | StrOutputParser()

Adding this chain to add_routes works perfectly with all the endpoints along with the per_req_config_modifier for passing on the user ID from the header:

async def per_req_config_modifier(config: Dict, request: Request) -> Dict:
    """
    Modify the config for each request
    """
    user = await get_current_active_user_from_request(request)
    config["configurable"] = {}
    # Attention: Make sure that the user ID is over-ridden for each request.
    # We should not be accepting a user ID from the user in this case!
    config["configurable"]["user_id"] = user["uid"]
    return config
   
add_routes(
    app,
    chain,
    per_req_config_modifier=per_req_config_modifier,
)