Implementing Real-Time Responses with LangChain and LLM
Tags: ai, openai, webdev, python
Published: 2023-12-24 06:44:03

I usually create chat tools for work, but I've always believed it's necessary to have knowledge about unmanned chatbots powered by AI. So recently, I've been experimenting with LangChain.

An article titled "Mitoyo City and Matsuo Lab's Half-Year Journey: Why They 'Gave Up' on ChatGPT for Business Efficiency Despite a 94% Accuracy Rate – Couldn't Trust AI for Garbage Disposal Guidance" mentioned an improvement where they enabled "real-time (streaming) display of responses." Inspired by this, I decided to implement a similar feature.I referenced the following course and added my own code to it.

https://www.udemy.com/share/109Hba3@j12oaDHOKBWzO1jzpItgQ-ib8CcNrAiGvJv3WSTkFO3UVtOX3yrZbqy5zIeo74d2/

## Implementation Overview
In this implementation, we use LangChain to process and display LLM-based responses in real-time. The main components include:

- `ChatOpenAI`: A class to utilize OpenAI's models.
- `ChatPromptTemplate`: A template for formatting chat inputs.
- `LLMChain`: A chain for processing using LLM.
- `StreamingHandler`: A callback handler for processing responses from LLM in real-time.
- Threads and queues: Elements for asynchronous processing and data exchange.

## Design
To achieve real-time display, the design includes the following elements:

1. Asynchronous Processing: Using Python threads, we execute LLM processes asynchronously. This ensures that the UI (not implemented in this case) does not freeze while waiting for responses from LLM, allowing other processes to continue.

2. Queues for Data Consistency: Using queues for data transfer between threads prevents conflicts when multiple threads access data simultaneously, maintaining data consistency. This is crucial in real-time systems to avoid data inconsistencies and unexpected errors, thereby enhancing the overall system reliability.

3. Streaming Processing: By obtaining and processing responses from LLM token by token, we provide real-time feedback to users.

## Implementation
I'll explain the code (available on [GitHub](https://github.com/suzuki0430/langchain-streaming-chain-test)). If you have obtained the `OPEN_API_KEY`, you can also check its behavior with Docker.

### Basic Setup
Import necessary libraries and call `load_dotenv()` to load environment variables. This safely manages external configuration information (`OPENAI_API_KEY`).

```python
import logging
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain
from langchain.callbacks.base import BaseCallbackHandler
from dotenv import load_dotenv
from queue import Queue
from threading import Thread

load_dotenv()
```

### StreamingHandler
The `StreamingHandler` class is a custom handler for processing LLM responses in real-time. It is called whenever new tokens are received from LLM or in case of an error.

- `on_llm_new_token`: Adds new tokens to the queue.
- `on_llm_end`: Notifies the queue when processing ends.
- `on_llm_error`: Logs errors and notifies the queue when processing ends.

```python
class StreamingHandler(BaseCallbackHandler):
    def __init__(self, queue):
        self.queue = queue

    def on_llm_new_token(self, token, **kwargs):
        self.queue.put(token)

    def on_llm_end(self, response, **kwargs):
        self.queue.put(None)

    def on_llm_error(self, error, **kwargs):
        logging.error(f"Error in LLM: {error}")
        self.queue.put(None)
```

### StreamingChain
The `StreamingChain` class is the main class for streaming data from LLM. It uses threads and queues to process LLM responses in real-time.

- `stream` method: Initiates LLM based on input and starts the result-generating process, which runs on a separate thread. The main thread continues to retrieve tokens from the queue.
- `cleanup` method: Waits for the thread to finish if it is still running after streaming ends.

```python
class StreamingChain:
    def __init__(self, llm, prompt):
        self.llm_chain = LLMChain(llm=llm, prompt=prompt)
        self.thread = None

    def stream(self, input):
        queue = Queue()
        handler = StreamingHandler(queue)

        def task():
            self.llm_chain(input, callbacks=[handler])

        self.thread = Thread(target=task)
        self.thread.start()

        try:
            while True:
                token = queue.get()
                if token is None:
                    break
                yield token
        finally:
            self.cleanup()

    def cleanup(self):
        if self.thread and self.thread.is_alive():
            self.thread.join()
```

### Usage Example
An example of using `StreamingChain` to obtain and display LLM responses in real-time based on user input. Here, we show streaming responses for the user input "Explain Pokémon in 100 characters."

```python
chain = StreamingChain(llm=chat, prompt=prompt)

for output in chain.stream(input={"content": "Explain Pokémon in 100 characters."}):
    print(output)
```

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/q1bhri5bur6fbb8584et.gif)