How to change from astream to astream_events

[levalencia]

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 for chunk in accumulate_and_yield(conversational_rag_chain.astream_({"input": prompt}, config={"configurable": {"session_id": "unbound"}})):

                yield chunk


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

            async for chunk in chunks:
                chunk_str = str(chunk)
                #logging.info(f"chunks coming in: {chunk_str}")

                # Detect context start
                if '{\'context\': [Document(metadata' in chunk_str:
                    in_context = True
                    start_index = chunk_str.index('{\'context\': [Document(metadata')
                    accumulated_context = chunk_str[start_index:]

                    # Check if the context ends in the same chunk
                    if ']}' in chunk_str[start_index:]:
                        end_index = chunk_str.index(']}', start_index) + 2
                        accumulated_context = chunk_str[start_index:end_index]
                        in_context = False

                        # Yield complete context
                        json_context = await convert_to_json(accumulated_context)
                        #logging.info(f"Complete context: {json_context.encode()}")                     
                        yield json_context.encode()
                        accumulated_context = ""
                    continue

                # Accumulate context
                if in_context:
                    accumulated_context += chunk_str

                    # Detect context end
                    if ']}' in chunk_str:
                        end_index = chunk_str.index(']}') + 2
                        accumulated_context = accumulated_context[:end_index]
                        in_context = False

                        # Yield complete context
                        json_context = json.dumps(accumulated_context)
                        #logging.info(f"Complete context: {json_context.encode()}")  
                        yield json_context.encode()
                        accumulated_context = ""
                    continue

                # Accumulate answers and detect "SOURCES:"
                if "{'answer':" in chunk_str:
                    answer_part = await parse_json_string(chunk_str)
                    if answer_part is None:
                        logging.error(f"Failed to parse JSON from chunk: {chunk_str}")
                        continue
                
                    if 'answer' not in answer_part:
                        logging.error(f"'answer' key not found in parsed JSON: {answer_part}")
                        continue
                    
                    answer_part = answer_part['answer']

                    # Handle potential partial detection of "SOURCES:"
                    partial_sources += answer_part
                    # logging.info(f"Partial Sources: {partial_sources}")
                    if 'SOURCES:' in partial_sources:
                        source_detected = True
                        parts = partial_sources.split('SOURCES:', 1)
                        accumulated_answer += parts[0]
                        sources_content = parts[1]
                        partial_sources = ""  # Reset partial sources tracker
                    else:
                        if source_detected:
                            sources_content += answer_part
                        else:
                            accumulated_answer += answer_part

                    # Yield accumulated answer if length >= 8 characters and no 'SOURCES:' detected yet
                    if len(accumulated_answer) >= 8 and not source_detected:
                        try:
                            json_answer = json.dumps({'answer': accumulated_answer})
                            #logging.info(f"Answer chunks: {json_answer.encode()}")
                            yield json_answer.encode()
                            accumulated_answer = ""
                        except Exception as e:
                            logging.error(f"Failed to process chunk: {chunk_str}")
                            logging.exception(e)
                            # Optionally yield an error or handle differently
                            continue

            # Yield any remaining accumulated answer that didn't reach 8 characters
            if accumulated_answer and not source_detected:
                try:
                    json_answer = json.dumps({'answer': accumulated_answer})
                    #logging.info(f"Final answer chunk: {json_answer.encode()}")
                    yield json_answer.encode()
                except Exception as e:
                    logging.error(f"Failed to process final accumulated answer: {accumulated_answer}")
                    logging.exception(e)

             # Yield the complete sources content if 'SOURCES:' was detected
            if source_detected:
                try:
                    json_sources = json.dumps({'SOURCES': sources_content})
                    #logging.info(f"Sources chunk: {json_sources.encode()}")
                    yield json_sources.encode()
                except Exception as e:
                    logging.error(f"Failed to process sources content: {sources_content}")
                    logging.exception(e)
    except Exception as e:
        logging.error(f"An error occurred in accumulate_and_yield: {e} {traceback.format_exc()}")
        yield f"An error occurred: {e} {traceback.format_exc()} ".encode()  
    finally:
        # Final yield to signal the end of the response
        yield b""

Description

I want to use SSE instead of streaming, how can I change my code from astream to astream_events

System Info

langchain_core: 0.2.21
langchain: 0.2.9
langchain_community: 0.2.7
langsmith: 0.1.92
langchain_openai: 0.1.17
langchain_text_splitters: 0.2.2
langchainhub: 0.1.20

[dosubot[bot]]

To change your code from using astream to astream_events to use SSE instead of streaming, you can follow this example:

Original code using astream:

async for chunk in accumulate_and_yield(conversational_rag_chain.astream_({"input": prompt}, config={"configurable": {"session_id": "unbound"}})):
    yield chunk

Modified code using astream_events:

events = []
async for event in accumulate_and_yield(conversational_rag_chain.astream_events({"input": prompt}, config={"configurable": {"session_id": "unbound"}}, version="v2")):
    events.append(event)
    kind = event["event"]
    if kind == "on_chat_model_stream":
        yield f"Chat model chunk: {repr(event['data']['chunk'].content)}".encode()
    if kind == "on_parser_stream":
        yield f"Parser chunk: {event['data']['chunk']}".encode()

This modification will allow you to stream results from intermediate steps using SSE. Note that if you are running Python <= 3.10, you may need to manually propagate the RunnableConfig object to the child runnable in async environments to ensure events are emitted correctly. For Python >= 3.11, this propagation happens automatically ^[1][2].

To continue talking to Dosu, mention @dosu.

[levalencia]

@dosu but for SSE I understand data should be returned with data: and then \n\n to end each event?

[dosubot[bot]]

To format the data for Server-Sent Events (SSE) with the data: prefix and \n\n suffix, you can modify the code to include these formatting requirements. Here's how you can adjust your code:

async for event in accumulate_and_yield(conversational_rag_chain.astream_events({"input": prompt}, config={"configurable": {"session_id": "unbound"}}, version="v2")):
    kind = event["event"]
    if kind == "on_chat_model_stream":
        yield f"data: Chat model chunk: {repr(event['data']['chunk'].content)}\n\n".encode()
    if kind == "on_parser_stream":
        yield f"data: Parser chunk: {event['data']['chunk']}\n\n".encode()

This modification ensures that each event is formatted correctly for SSE by adding the data: prefix and \n\n suffix to each chunk of data before yielding it.

To continue talking to Dosu, mention @dosu.

[levalencia]

@dosu what I want to differentiate between the context chunk and the answer chunks.

[dosubot[bot]]

To differentiate between the context chunk and the answer chunks in your SSE implementation, you can use the keys "context" and "answer" respectively. Here's how you can modify your code to include these keys and format the data for SSE:

async for event in accumulate_and_yield(conversational_rag_chain.astream_events({"input": prompt}, config={"configurable": {"session_id": "unbound"}}, version="v2")):
    kind = event["event"]
    if kind == "on_chat_model_stream":
        yield f"data: {{'answer': {repr(event['data']['chunk'].content)}}}\n\n".encode()
    elif kind == "on_parser_stream":
        yield f"data: {{'context': {event['data']['chunk']}}}\n\n".encode()

In this code:

• The answer chunks are yielded with the key "answer".
• The context chunks are yielded with the key "context".

This ensures that each chunk of data is correctly identified and formatted for SSE with the data: prefix and \n\n suffix ^[1][2].

To continue talking to Dosu, mention @dosu.