Returning a chain with a custom runnable instead of just a runnable

[rarchitgupta]

I am trying to use the per_req_config modifier to pass in a user_id as well as an openai_api_key to my chain

It looks 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]
    openai_api_key: Optional[str]
    vectorstore: Optional[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):
        self.vectorstore = get_vectorstore(self.openai_api_key)
        retriever = self.vectorstore.similarity_search(
            query, k=5, 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

This works with the chain:

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

prompt = hub.pull("rlm/rag-prompt")
llm = get_llm(assume_api_key)

chain = per_user_retriever | prompt | llm | StrOutputParser()

However I am unable to extend the chain like so:

def get_rag_chain(self, query):
        self.vectorstore = get_vectorstore(self.openai_api_key)
        retriever = self.vectorstore.similarity_search(
            query, k=5, filter={"user_id": {"$eq": self.user_id}}
        )

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

        prompt = hub.pull("rlm/rag-prompt")
        llm = get_llm(self.openai_api_key)

        return (
            {
                "context": format_docs(retriever),
                "question": RunnablePassthrough(),
            }
            | prompt
            | llm
        )

hoping to use it as:

chain = per_user_retriever | StrOutputParser()

Instead I get the error:

TypeError: Expected a Runnable, callable or dict.Instead got an unsupported type: <class 'str'>

Any help with this would be appreciated