How to properly get answer chunks vs context chunks

[levalencia]

I have the following code which uses openan and streaming responses. However sometimes the context chunk comes in the answer chunks.

Any idea why?

async def regulatory_dossier_gpt_chain_lcel(prompt, username, inputs, outputs,access_token_keyvault, access_token_openai, access_token_search, access_token_storage, session_id, metadata_filters):
    async def llm_thread_gpt_lcel(prompt, username, inputs, outputs, metadata_filters):
        start_time = time.time()
        try:
            dep_start_time = time.time()
            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)
            SHD_STORAGE_ACCOUNT_NAME = config.SHD_STORAGE_ACCOUNT_NAME(access_token_keyvault)
         
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "fetch_config_execution_time", "session_id": session_id})

            dep_start_time = time.time()
            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)   
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "azurechatopenai_initialization_time", "session_id": session_id})

            dep_start_time = time.time()
            embeddings = setup_embeddings(access_token_keyvault, access_token_openai)
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "embeddings_setup_time", "session_id": session_id})

            dep_start_time = time.time()
            vector_store = setup_vector_store(embeddings,"index-regulatorydossier", access_token_keyvault, access_token_search)
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "vector_store_setup_time", "session_id": session_id})


            dep_start_time = time.time()
            retriever = ArgenxUserRetrieverRegulatoryDossierLCEL(vectorstore=vector_store, username=username, metadata_filters=metadata_filters)
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "retriever_setup_time", "session_id": session_id})

            
            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} Please answer in HTML format except for the SOURCES: section and make sure to always start your citations in the answer with the number [1] than [2] and so on"),
                ]
            )
            dep_start_time = time.time()
            history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt) 
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "history_aware_retriever_setup_time", "session_id": session_id})

            dep_start_time = time.time()           
            ANSWER_PROMPT_REG_DOSSIER = read_text_from_blob(
                SHD_STORAGE_ACCOUNT_NAME, 
                BLOB_CONTAINER_NAME, 
                REGDOSSIER_PROMPT_BLOB_NAME,
                access_token_storage
            )
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "read_text_from_blob_regdossierprompt", "app": "regulatory", "step": "history_aware_retriever_setup_time", "session_id": session_id})


            dep_start_time = time.time()
            qa_prompt = ChatPromptTemplate.from_messages(
                [
                    ("system", ANSWER_PROMPT_REG_DOSSIER),
                    MessagesPlaceholder("chat_history"),
                    ("human", "{input} Please answer in HTML format except for the SOURCES: section and make sure to always start your citations in the answer with the number [1] than [2] and so on"),
                ]
            )
            
             # Define your legacy DOCUMENT_PROMPT
            document_prompt = PromptTemplate(
                input_variables=["page_content", "source"],
                template="Content: {page_content}\nSource: {source}"
            )
            
            question_answer_chain = create_stuff_documents_chain(llm, qa_prompt,document_prompt=document_prompt, document_variable_name="context" )  | JsonOutputParser()
            dep_end_time = time.time()
            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "question_answer_chain_setup_time", "session_id": session_id})


            dep_start_time = time.time()   
            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",
            )
            dep_end_time = time.time()

            meter_function_execution_time.record(dep_end_time - dep_start_time, {"function_name": "regulatory_dossier_gpt_chain_lcel", "app": "regulatory", "step": "rag_chain_setup_time", "session_id": session_id})

            context_check_count = 0
            max_context_checks = 10
            context_detected = False
            context_yielded = False
            accumulated_chunk = ""
            sourcing_detected = False
            ttft_start_time = time.time()
             
            async for event in conversational_rag_chain.astream_events({"input": prompt}, config={"configurable": {"session_id": "unbound"}}, version="v2"):
                kind = event["event"]

                if not context_detected:
                    if kind == "on_chain_end":
                        output_data = event['data'].get('output')
                        from langchain.schema import Document

                        if isinstance(output_data, list):
                            context_list = []  # Initialize an empty list to collect context data
                            for document in output_data:
                                if isinstance(document, Document):
                                    context_data = document.metadata  # Collect the metadata
                                    context_list.append(context_data)  # Append to the list

                            # After the loop, check if we have accumulated context
                            if context_list:  # Ensure there's something to yield
                                ttft_end_time = time.time()
                                ttft_duration = ttft_end_time - ttft_start_time
                                meter_time_first_token.record(ttft_duration, {"app": "regulatory", "session_id": session_id})

                                context_detected = True
                                context_yielded = True  # Set the flag to indicate that context has been yielded
                                yield f'data: {{"context": "{context_list}"}} \n\n'.encode()
                        
                        context_check_count += 1
                        
                        if not context_yielded and context_check_count >= max_context_checks:
                            # If we've checked max_context_checks times and still no context, proceed without context
                            context_detected = True
                            context_yielded = True  # Ensure we don't check for context again
                            yield f'data: {{"context": "[]"}} \n\n'.encode()

                else:  # context_detected is True, now process the answer
                    if kind == "on_parser_stream":
                        # Replace newline characters in the chunk value
                        chunk_value = event["data"]["chunk"].replace("\n", "\\n")
                        
                        # Accumulate the chunk values
                        accumulated_chunk += chunk_value
                        
                        if len(accumulated_chunk) < 16:
                            continue
                        
                        if sourcing_detected:
                            # If SOURCES: has already been detected, accumulate and output the sources content only
                            if "|||" in accumulated_chunk:
                                pipes_index = accumulated_chunk.index("|||")
                                content_after_pipes = accumulated_chunk[pipes_index+3:]                            
                                yield f'sources: {content_after_pipes} [END]'.encode()
                            else:
                                yield f'sources: {accumulated_chunk} [END]'.encode()
                            
                            # Clear the accumulator after yielding
                            accumulated_chunk = ""
                        else:                        
                            # Output content in chunks of 16 characters minimum
                            while len(accumulated_chunk) > 16:
                                # Find the next split point, ensuring \n is not broken
                                split_index = min(16, len(accumulated_chunk))
                                
                                if "|||" in accumulated_chunk:
                                    pipes_index = accumulated_chunk.index("|||")
                                    # Separate the content before and after `|||`
                                    content_before_pipes = accumulated_chunk[:pipes_index]
                                    yield f'answer: {content_before_pipes} [END]'.encode()
                                    yield 'endanswer: endanswer [END]'.encode()
                                    sourcing_detected = True
                                    break
                                else:                                
                                    # Check if the chunk ends with any of the special sequences
                                    if (accumulated_chunk[:split_index].endswith("\\") or
                                        accumulated_chunk[:split_index].endswith("\\n") or
                                        accumulated_chunk[:split_index].endswith("\\\\") or
                                        accumulated_chunk[:split_index].endswith("\\n\\") or
                                        accumulated_chunk[:split_index].endswith("\\n\\n")):                                        
                                        # If so, we should accumulate the next chunk before yielding
                                        # Extend the split_index to the next part of the chunk
                                        if len(accumulated_chunk) > split_index:
                                            split_index += 1  # Expand to include next character
                                            # If it's still a special sequence, include further characters
                                            while (split_index < len(accumulated_chunk) and 
                                                (accumulated_chunk[:split_index].endswith("\\") or
                                                    accumulated_chunk[:split_index].endswith("\\n") or
                                                    accumulated_chunk[:split_index].endswith("\\\\"))):
                                                split_index += 1                                                
                                    if  not accumulated_chunk.endswith("|") and not accumulated_chunk.endswith("||"):
                                        chunk_to_yield = accumulated_chunk[:split_index]
                                        yield f'answer: {chunk_to_yield} [END]'.encode()
                                        # Update accumulated_chunk by removing the processed chunk
                                        accumulated_chunk = accumulated_chunk[split_index:]
                                    else:
                                        logger.info('chunk ends with | or ||, accumulate it with next')
                                        break

            # Final check to yield any remaining chunk
            if accumulated_chunk:
                if sourcing_detected:
                    yield f'sources: {accumulated_chunk} [END] '.encode()
                    yield 'endsources: endsources [END]'.encode()
                else:
                    #logger.info(f'data: {{"sources": "{accumulated_chunk}"}} \n\n'.encode())
                    yield f'answer: {accumulated_chunk} [END] '.encode()
                    yield 'endanswer: endanswer [END]'.encode()
            # Ensure the final markers are yielded, regardless of accumulated_chunk
            else:
                if sourcing_detected:
                    # If SOURCES detected, yield endsources marker
                    yield 'endsources: endsources [END]'.encode()
                else:
                    # If no SOURCES detected, yield endanswer marker
                    yield 'endanswer: endanswer [END]'.encode()
        except Exception as e:
            logger.error(f"An error occurred in regulatory_dossier_gpt_chain_lcel: {e}")
            yield f"error: {e}".encode()         
            
    try:
        async for chunk in llm_thread_gpt_lcel(prompt, username, inputs, outputs, metadata_filters):
            yield chunk
    except Exception as e:
        logger.error(f"An error occurred in regulatory_dossier_gpt_chain_lcel: {e}")
        yield f"error: {e}".encode()

[dosubot[bot]]

I found a similar discussion that might be relevant to your issue: Answer chunks not coming in with JsonOutputParser ^[1]. This discussion addresses issues with chunked responses and might provide insights into handling context and answer chunks separately.

In your code, the issue of context chunks appearing in answer chunks could be due to how the streaming and chunk accumulation logic is implemented. Here are a few suggestions to help you troubleshoot and potentially resolve the issue:

1. Chunk Accumulation Logic: Ensure that the logic for accumulating and yielding chunks is correctly distinguishing between context and answer chunks. You might need to refine the conditions under which chunks are accumulated and yielded, especially around the detection of special sequences or markers like |||.

2. Event Handling: Check the handling of different event types (on_chain_end, on_parser_stream, etc.) to ensure that context and answer data are processed separately. Make sure that the context is fully processed and yielded before starting to process answer chunks.

3. Markers and Delimiters: Use clear and distinct markers or delimiters to separate context from answer chunks. This can help in correctly identifying and processing each type of chunk.

4. Debugging and Logging: Add more detailed logging to track the flow of data and identify where context might be incorrectly included in answer chunks. This can help pinpoint the exact step where the issue occurs.

By refining these areas, you should be able to better manage the separation of context and answer chunks in your streaming responses.

To continue talking to Dosu, mention @dosu.

[dosubot[bot]]

To return answer chunks and sources chunks effectively, you can use specific markers or delimiters to distinguish between them. In your code, you can use markers like "SOURCES:", "Sources:", "source:", and "Source:" to identify the beginning of the sources chunk. This approach ensures that the sources chunk is clearly delineated from the answer chunk ^[1][2].

Regarding the events, here are some key ones you might encounter:

1. on_chain_end: Typically used to handle the end of a chain process, which might include context or metadata.
2. on_parser_stream: Used for streaming parsed data, where you can accumulate and yield answer chunks.
3. on_chat_model_stream: Handles streaming of chat model output, which can include answer chunks.

By focusing on these markers and events, you can better manage the separation and processing of answer and sources chunks.

To continue talking to Dosu, mention @dosu.

[levalencia]

@dosu, based on my code, how do I correctly get the context chunks?

[dosubot[bot]]

To correctly get the context chunks in your code, you should focus on the logic within the llm_thread_gpt_lcel function, particularly how you handle the on_chain_end event. Here's a refined approach:

1. Event Handling: Ensure that the on_chain_end event is specifically used to detect and process context chunks. This event is where you check if the output data contains context information.

2. Context Detection: Use a clear condition to identify context chunks. In your code, you are checking if the output_data is a list of Document objects and then extracting the metadata as context. Ensure this logic is robust and correctly identifies context data.

3. Yielding Context: Once context is detected, yield it immediately and set flags (context_detected and context_yielded) to prevent further checks for context in subsequent events. This ensures that context is processed and output before any answer chunks.

4. Markers and Delimiters: If needed, use specific markers or delimiters to clearly separate context from other types of data. This can help in distinguishing context chunks from answer or source chunks.

By focusing on these steps, you can ensure that context chunks are correctly identified and processed separately from other types of chunks in your streaming responses.

To continue talking to Dosu, mention @dosu.

[levalencia]

@dosu what about if use astream instead of astream_events?

[dosubot[bot]]

Could you provide more details on how you plan to use astream instead of astream_events? Specifically, are there any particular changes or improvements you are looking to achieve by switching to astream? Additionally, understanding how you currently handle events and chunks in your code will help in providing a more tailored solution.

To continue talking to Dosu, mention @dosu.