Add Streaming Support to RAG Pipeline

Summary

This PR introduces token-level streaming to the RAG pipeline, enhancing interactivity and responsiveness. The new RagStreamHandler class manages real-time response streaming and integrates into RagPipelineWrapper.

Key Changes

• Streaming Handler:
  • Added RagStreamHandler for real-time token streaming with methods to start, stop, and yield response chunks.
• Pipeline Integration:
  • Introduced stream_query in RagPipelineWrapper to handle streaming queries.
  • Integrated streaming callbacks into _add_llm.
  • Updated run to switch between streaming and non-streaming modes based on the stream setting.
• Configuration:
  • Added STREAM_TIMEOUT to manage streaming timeout.
  • Enabled stream in settings.py:

    "stream": True

Testing

• Verified token-level streaming works correctly.
• Confirmed compatibility with non-streaming mode.
• Ensured proper thread and queue management.

Impact

• Improved Responsiveness: Faster, interactive responses with streaming.
• No Regression: Non-streaming functionality remains intact.

How to Use

1. Enable streaming in settings.py:

   "stream": True

2. Use the stream_query method in RagPipelineWrapper to enable streaming queries.
3. Handle and print streaming responses using the main.py script:

   result_generator = execute_rag_query(args.query, **custom_settings)
   for chunk in result_generator:
       decoded_chunk = chunk.decode("utf-8")
       print(decoded_chunk, end="", flush=True)

Notes for Reviewers

• Review RagStreamHandler for robustness and edge-case handling.
• Verify pipeline integration ensures consistent streaming behavior.