Return sources with metadata using Langserv

[kcole93]

I'm trying to wrap my head around how to modify the schema of the output returned by the invoke endpoint exposed by LangServe.

I believe this is the relevant portion of my chain.py file inside my LangServe app:

embeddings = BedrockEmbeddings(
    model_id="amazon.titan-embed-text-v1", region_name="us-east-1"
)

vectorstore = Chroma.from_documents(documents=data, embedding=embeddings)

# Get retriever from vectorstore
retriever = vectorstore.as_retriever()

prompt_template = """Use the following pieces of context to answer the question at the end.
If you don't know the answer from the provided context, just say that your training materials don't include this information, don't try to make up an answer.
Keep the answer as concise as possible.

{context}

Question: {question}

Helpful Answer:"""
custom_rag_prompt = PromptTemplate.from_template(prompt_template)

llm = Bedrock(model_id="amazon.titan-text-express-v1", model_kwargs={"maxTokenCount": 4000})

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


rag_chain_from_docs = (
    RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
    | custom_rag_prompt
    | llm
    | StrOutputParser()
)


rag_chain_with_source = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough()}
).assign(answer=rag_chain_from_docs)

chain = rag_chain_with_source

The response from returned from this chain by the API has the shape of:

{
  "output": {
     "question": "string",
     "answer": "string",
     },
     "callback_events": [],
     "metadata": {
         "run_id": "string"
       }
}

I want to return not just the question and answer, but also context -- more specifically, I want to process the context returned to extract values from the source documents' metadata (urls pointing to their origin), and return this as either output or metadata for each call to the invoke endpoint. So far, I haven't been able to understand how to modify the schema of the output.

I'd appreciate any tips on where changes can be made to the output's schema! Thanks!

[dosubot[bot]]

🤖

Hey @kcole93! Good to see you back here. Hope you're doing well and your coding adventures are exciting as always. Let's dive into this new issue!

To modify the output schema of the invoke endpoint in LangServe, you can create a custom output parser. This parser can be used to format the output of your chain in the desired schema.

In your case, you want to include the context and the URLs pointing to the origin of the source documents. To achieve this, you can modify your format_docs function to return not just the formatted documents, but also a list of URLs. You can then include this list in the output of your chain.

Here's an example of how you can modify your format_docs function:

def format_docs(docs):
    formatted_docs = "\n\n".join(doc.page_content for doc in docs)
    urls = [doc.url for doc in docs]  # assuming each doc has a 'url' attribute
    return formatted_docs, urls

Then, you can modify your rag_chain_from_docs to include the URLs in the output:

rag_chain_from_docs = (
    RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])[0]))
    | custom_rag_prompt
    | llm
    | StrOutputParser()
    | RunnablePassthrough.assign(urls=(lambda x: format_docs(x["context"])[1]))
)

In this code, RunnablePassthrough.assign(urls=(lambda x: format_docs(x["context"])[1])) is used to add the URLs to the output. The lambda function lambda x: format_docs(x["context"])[1] extracts the URLs from the context.

Please note that this is a simplified example and you might need to adjust it according to your specific requirements and the structure of your documents.

As for including the context in the output, you can do this by modifying the RunnableParallel in your rag_chain_with_source:

rag_chain_with_source = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough(), "urls": RunnablePassthrough()}
).assign(answer=rag_chain_from_docs)

In this code, RunnablePassthrough() is used to pass the context and URLs to the output.

I hope this helps! If you have any further questions, feel free to ask.

For more information on how to use output parsers in LangChain, you can refer to the LangChain documentation.

Sources

• libs/langchain/langchain/output_parsers/structured.py
• libs/core/langchain_core/runnables/base.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[Chengdyc]

Hi, I'm running into the same problem, did you figure out a solution? the solution proposed by dosubot doesn't help because the chain is already returning the context, but the langserve APIHandler is not passing the context along with the answer.

[alice17]

I'm having the same problem

[eyurtsev]

What's the output from chain if you do .invoke without langserve?

[alice17]

In the end I solved by looking at this issue. It would be nice to have this example in the documentation. Now I have to find a way to return only urls.

[alice17]

Found a way:

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


def get_urls(docs):
    return list(set([x.metadata["source"] for x in docs["context"]]))


def get_chain(retriever):
    rag_chain_from_docs = (
            {
                "question": lambda x: x["question"],
                "context": lambda x: format_docs(x["context"]),
            }
            | prompt
            | llm
            | StrOutputParser()
    )

    retrieve_docs = (lambda x: x["question"]) | retriever

    return (RunnablePassthrough.assign(context=retrieve_docs)
            .assign(answer=rag_chain_from_docs)
            .assign(sources=(lambda x: get_urls(x))))

Now /invoke output looks like this:

{
	"output": {
		"question": "question",
		"answer": "anwer",
		"sources": [
			"www.google.com", "...."]
	},
	"metadata": {
		"run_id": "",
		"feedback_tokens": []
	}
}

The only problem is that now the answers in the playground are totally messed up.