Answer chunks not coming in with JsonOutputParser

[Pieter-JanCnuddeE61]

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

async def msl_gpt_chain_lcel(prompt, username, inputs, outputs,access_token_keyvault, access_token_openai, access_token_search):
    async def llm_thread_gpt_lcel(prompt, username, inputs, outputs):
        
        try:
            SHD_AZURE_OPENAI_ENDPOINT = config.SHD_AZURE_OPENAI_ENDPOINT(access_token_keyvault)
            SHD_OPENAI_DEPLOYMENT_NAME = config.SHD_OPENAI_DEPLOYMENT_NAME(access_token_keyvault)
            SHD_OPENAI_GPT_MODEL_NAME = config.SHD_OPENAI_GPT_MODEL_NAME(access_token_keyvault)
   
            llm = AzureChatOpenAI(
                azure_endpoint=SHD_AZURE_OPENAI_ENDPOINT,
                openai_api_version="2023-03-15-preview",
                deployment_name=SHD_OPENAI_DEPLOYMENT_NAME,
                azure_ad_token = access_token_openai,
                openai_api_type="Azure",
                model_name=SHD_OPENAI_GPT_MODEL_NAME,
                streaming=True,  # Set ChainStreamHandler as callback
                temperature=0)     

            embeddings = setup_embeddings(access_token_keyvault, access_token_openai)
            vector_store = setup_vector_store(embeddings, access_token_keyvault, access_token_search)
            retriever = ArgenxUserScoringRetrieverMSLLCEL(vectorstore=vector_store, username=username)
            
            demo_ephemeral_chat_history = ChatMessageHistory()
            for i in range(len(inputs)):
                demo_ephemeral_chat_history.add_user_message(inputs[i])
                demo_ephemeral_chat_history.add_ai_message(outputs[i])
                
            contextualize_q_system_prompt = (
                "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."
            )
            contextualize_q_prompt = ChatPromptTemplate.from_messages(
                [
                    ("system", contextualize_q_system_prompt),
                    MessagesPlaceholder("chat_history"),
                    ("human", "{input}"),
                ]
            )
            history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt) 
            qa_prompt = ChatPromptTemplate.from_messages(
                [
                    ("system", ANSWER_PROMPT_MSL),
                    MessagesPlaceholder("chat_history"),
                    ("human", "{input}"),
                ]
            )

             # Define your legacy DOCUMENT_PROMPT
            document_prompt = PromptTemplate(
                input_variables=["page_content", "source", "document_type", "year", "document_region"],
                template="START DOCUMENT {source}:\n '{page_content}' \n METADATA ABOUT DOCUMENT {source}: 1. Document type: {document_type} - Year: {year} \n 2. Source: {source} \n 3. Region: {document_region} \n END DOCUMENT \n"
            )

            question_answer_chain = create_stuff_documents_chain(llm, qa_prompt, document_prompt=document_prompt, document_variable_name="context")  | JsonOutputParser()
                
            rag_chain_from_docs = create_retrieval_chain(history_aware_retriever, question_answer_chain) 
            conversational_rag_chain = RunnableWithMessageHistory(
                rag_chain_from_docs,
                lambda session_id: demo_ephemeral_chat_history, 
                input_messages_key="input",
                history_messages_key="chat_history",
                output_messages_key="answer",
            )

            async for chunk in accumulate_and_yield_msl(conversational_rag_chain.astream({"input": prompt}, config={"configurable": {"session_id": "unbound"}})):
                yield chunk
        except Exception as e:
            log_exception(tc,e)
            logging.error(f"An error occurred in llm_thread_gpt_lcel: {e}")
            yield f"An error occurred: {e}".encode()
            
    try:
        async for chunk in llm_thread_gpt_lcel(prompt, username, inputs, outputs):
            yield chunk
    except Exception as e:
        log_exception(tc,e)
        logging.error(f"An error occurred in msl_gpt_chain_lcel: {e}")
        yield f"An error occurred: {e}".encode()

Description

Before JsonOutputParser I was using the StrOutputParser and I could see all the chunks coming in for context, input and answer. However, since I changed to the JsonOutputParser, only the input and context chunk are coming in. The chunk following the context chunk is None.

async def accumulate_and_yield_msl(chunks):
    accumulated_answer = ""
    accumulated_context = ""
    in_context = False
    source_detected = False
    sources_content = ""
    partial_sources = ""

    async for chunk in chunks:
        # Skip processing if chunk is None
        if chunk is None:
            continue

        chunk_str = str(chunk)

System Info

System Information

OS: Windows
OS Version: 10.0.22621
Python Version: 3.11.9 (tags/v3.11.9:de54cf5, Apr 2 2024, 10:12:12) [MSC v.1938 64 bit (AMD64)]

Package Information

langchain_core: 0.2.25
langchain: 0.2.11
langchain_community: 0.2.10
langsmith: 0.1.94
langchain_openai: 0.1.19
langchain_text_splitters: 0.2.2
langchainhub: 0.1.20

[dosubot[bot]]

Hey there, @PjCnuddelement! I'm a bot here to help you with bugs, questions, and more while you wait for a human maintainer. If you need assistance, feel free to ask!

To resolve the issue where the chunk following the context chunk is None when using JsonOutputParser, ensure that the additional_kwargs and tool_calls are correctly populated in the AIMessageChunk instances. Here is a relevant part of the code that ensures the tool_call_chunks are correctly populated:

STREAMED_MESSAGES_WITH_TOOL_CALLS = []
for message in STREAMED_MESSAGES:
    if message.additional_kwargs:
        STREAMED_MESSAGES_WITH_TOOL_CALLS.append(
            AIMessageChunk(
                content=message.content,
                additional_kwargs=message.additional_kwargs,
                tool_call_chunks=[
                    ToolCallChunk(
                        name=chunk["function"].get("name"),
                        args=chunk["function"].get("arguments"),
                        id=chunk.get("id"),
                        index=chunk["index"],
                    )
                    for chunk in message.additional_kwargs["tool_calls"]
                ],
            )
        )
    else:
        STREAMED_MESSAGES_WITH_TOOL_CALLS.append(message)

Additionally, ensure that your JsonOutputParser is correctly handling partial JSON objects. The parse_result method in JsonOutputParser should handle partial JSON objects and return None if the JSON is incomplete or invalid:

class JsonOutputParser(BaseCumulativeTransformOutputParser[Any]):
    def parse_result(self, result: List[Generation], *, partial: bool = False) -> Any:
        text = result[0].text
        text = text.strip()
        if partial:
            try:
                return parse_json_markdown(text)
            except JSONDecodeError:
                return None
        else:
            try:
                return parse_json_markdown(text)
            except JSONDecodeError as e:
                msg = f"Invalid json output: {text}"
                raise OutputParserException(msg, llm_output=text) from e

Ensure that your STREAMED_MESSAGES list is correctly populated with additional_kwargs and tool_calls to avoid None values in the chunks, and that the JsonOutputParser is correctly handling partial JSON objects.

For further robustness, refer to the test cases provided in the LangChain library, which include various scenarios for parsing JSON, including partial JSON objects and JSON within code blocks ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[levalencia]

@dosu this answer is so wrong!