How to assign a specific llm field to an output key

[santiagovasquez1]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import os
from typing import Any, Dict
import logging
from Helpers.set_gpts__helper import set_gpt_embeddings, set_gpt_llm
from azure.functions import HttpRequest, HttpResponse
from langchain_community.vectorstores.azuresearch import AzureSearch
from Helpers.set_env_helper import set_env
from langchain_core.language_models import BaseLanguageModel
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableSerializable, RunnablePassthrough
from langchain.chains import create_retrieval_chain, LLMChain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.output_parsers import StrOutputParser
from langchain.chains import create_history_aware_retriever, SequentialChain
from langchain.output_parsers import JsonOutputKeyToolsParser

set_env()
history = ChatMessageHistory()
llm = set_gpt_llm(
    deployment=os.getenv("AZURE_CHAT_DEPLOYMENT"),
    api_version=os.getenv("AZURE_OPENAI_VERSION")
)

embeddings = set_gpt_embeddings(
    deployment=os.getenv("AZURE_EMBEDDINGS_DEPLOYMENT"),
    api_version=os.getenv("AZURE_OPENAI_VERSION")
)
index_name = os.getenv("AZURE_SEARCH_INDEX_NAME", "rag-azure-search")

vector_store_connection = AzureSearch(azure_search_endpoint=os.environ["AZURE_SEARCH_ENDPOINT"],
                                      azure_search_key=os.environ["AZURE_SEARCH_KEY"],
                                      index_name=index_name,
                                      embedding_function=embeddings.embed_query)

retriever = vector_store_connection.as_retriever(
    search_type="similarity", k=3, search_kwargs={"score_threshold": 0.7})


def main(req: HttpRequest) -> HttpResponse:

    history_aware_retriever_chain = create_history_aware_retriever(
        llm, retriever, contextualize_system_prompt()
    )
    combine_docs_chain = create_stuff_documents_chain(llm, qa_system_prompt())

    summary_chain = summarize_system_chain(llm)

    rag_chain: RunnableSerializable[Dict, Any] = create_retrieval_chain(
        history_aware_retriever_chain, combine_docs_chain)

    query = req.params.get('name')
    if not query:
        try:
            req_body = req.get_json()
        except ValueError:
            pass
        else:
            query = req_body.get('name')

    if query:
        response = rag_chain.invoke(
            {"input": query, "chat_history": history.messages})
        summary_response = summary_chain.invoke(response["answer"])
        history.add_user_message(query)
        history.add_ai_message(summary_response)
        logging.info(response["system_message"]["answer"])
        return HttpResponse(response["system_message"]["answer"])
    else:
        return HttpResponse(
            "This HTTP triggered function executed successfully. Pass a name in the query string or in the request body for a personalized response.",
            status_code=200
        )


def summarize_system_chain(llm: BaseLanguageModel):
    template = "Create a summary with the main ideas of the following text:\n{system_message} Which was generated by a cognitive assistant."
    prompt = ChatPromptTemplate.from_template(template)
    chain = prompt | llm | StrOutputParser()
    return chain


def contextualize_system_prompt():
    template = """Given a chat history and the latest user question \
    which might reference context in the chat history, formulate a standalone question \
    which can be understood without the chat history. Do NOT answer the question, \
    just reformulate it if needed and otherwise return it as is. """

    return set_prompt(template)


def qa_system_prompt():
    template = """You are an assistant for question-answering tasks. \
    Use the following pieces of retrieved context to answer the question. \
    If you don't know the answer, just say that you don't know. \
    keep the answer concise.\

    {context}"""
    return set_prompt(template)


def set_prompt(template: str):
    return ChatPromptTemplate.from_messages(
        [
            ("system", template),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )


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

Description

I want to create a sequential chain in which first invoke rag_chain and then pass to summary_chain, how can I pass the rag_chain[answer] field as an inpute of summary_chain, I have tried with

full_chain = RunnablePassthrough().assign(
        system_message=rag_chain).assign(summary=summary_chain)

response = full_chain.invoke(
            {"input": query, "chat_history": history.messages})

but when I do it in the output field system_message comes all the rag_chain dictionary and not only the answer

System Info

langchain==0.2.14
langchain-cli==0.0.30
langchain-community==0.2.12
langchain-core==0.2.35
langchain-experimental==0.0.64
langchain-openai==0.1.22
langchain-text-splitters==0.2.2

[dosubot[bot]]

Hello there, @santiagovasquez1! 👋 I'm here to assist you with any bugs, questions, or contributions you might have. I'm a bot designed to help you while waiting for a human maintainer.

To create a sequential chain where you first invoke rag_chain and then pass the rag_chain["answer"] field as an input to summary_chain, you can use the RunnablePassthrough.assign method to extract and pass the specific field. Here's how you can achieve this:

1. Define the rag_chain to perform the initial question-answering task.
2. Extract the answer from the rag_chain output.
3. Pass the extracted answer to the summary_chain for summarization.

Here is an example of how you can achieve this:

from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# Define the RAG chain
RAG_TEMPLATE = """
You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.

<context>
{context}
</context>

Answer the following question:

{question}"""

rag_prompt = ChatPromptTemplate.from_template(RAG_TEMPLATE)
llm = ChatOpenAI(model="gpt-4")

rag_chain = (
    RunnablePassthrough.assign(context=lambda input: format_docs(input["context"]))
    | rag_prompt
    | llm
    | StrOutputParser()
)

# Define the summary chain
summary_prompt = ChatPromptTemplate.from_template(
    "Summarize the following answer: {answer}"
)

summary_chain = (
    RunnablePassthrough.assign(answer=lambda input: input["answer"])
    | summary_prompt
    | llm
    | StrOutputParser()
)

# Combine the chains
combined_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | rag_chain
    | RunnablePassthrough.assign(answer=lambda output: output["answer"])
    | summary_chain
)

# Invoke the combined chain
question = "What are the approaches to Task Decomposition?"
result = combined_chain.invoke(question)
print(result)

In this example:

• rag_chain is responsible for generating the initial answer based on the retrieved context.
• The answer from rag_chain is then passed to summary_chain for summarization.
• The combined_chain sequentially invokes rag_chain and then summary_chain using the RunnablePassthrough.assign method to pass the necessary fields between the chains.

This approach leverages the concepts of Retrieval Augmented Generation (RAG) as described in the LangChain documentation, where the retrieval and generation steps are combined to answer questions based on specific source information ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[santiagovasquez1]

@dosu

in the example you propose, how can I get the documents field of the result to see the documents on which the answer was based and also how can I see the message generated by the rag_chain?

[dosubot[bot]]

To get the documents field of the result to see the documents on which the answer was based and also see the message generated by the rag_chain, you can use the following code snippet:

response = rag_chain_with_source.invoke({"question": "What is Task Decomposition?"})
print(response["answer"])

for document in response["documents"]:
    print(document["page_content"])
    print()

This code will print the generated answer and the content of the documents that were retrieved and used to generate the answer. The documents field in the response contains the sources that the LLM used in generating the response ^[1][2][3].

To continue talking to Dosu, mention @dosu.