LLM Exploration: Adding Continuous Conversation Functionality to ChatGLM2's gRPC Backend

Preface

Previously, when I worked on AIHub, I integrated open-source large language models like ChatGLM via gRPC. For the large language model part, I developed the StarAI framework, which is essentially a simplified version of langchain. It allows for relatively convenient combination and use of various large language models and their supporting tools.

The main approach still follows the OpenAI interface pattern to reduce learning costs. However, to speed up development earlier, I only implemented a simple gRPC interface and hadn't yet completed the multi-turn dialogue feature. This time, I'll be refining and completing that functionality.

Overview

The system is divided into two parts: the LLM backend and the client. The LLM backend uses gRPC to provide interfaces, while the client is the AIHub I developed with Blazor.

So, this time involves modifications in the following areas:

• proto
• Client – C# code
• AIHub pages – Blazor razor code
• gRPC server – Python code

Modifying the Proto

Let's modify the proto file.

syntax = "proto3";

import "google/protobuf/wrappers.proto";

option csharp_namespace = "AIHub.RPC";

package aihub;

service ChatHub {
  rpc Chat (ChatRequest) returns (ChatReply);
  rpc StreamingChat (ChatRequest) returns (stream ChatReply);
}

message ChatRequest {
  string prompt = 1;
  repeated Message history = 2;
  int32 max_length = 3;
  float top_p = 4;
  float temperature = 5;
}

message Message {
  string role = 1;
  string content = 2;
}

message ChatReply {
  string response = 1;
}

Added the Message type, and introduced the history field in the ChatRequest chat request to serve as conversation history.

Modifying the C# gRPC Client Code

After writing the proto above and compiling the project, the C# client code will be regenerated. Now let's modify our calling code.

You can see that ChatRequest now has a history property of type RepeatedField<Message>. This property is read-only, so when chatting, you can only add conversation history by using the add method.

For convenience, I encapsulated the following method to create a ChatRequest object.

private ChatRequest GetRequest(string prompt, List<Message>? history = null) {
  var request = new ChatRequest {
    Prompt = prompt,
    MaxLength = 2048,
    TopP = 0.75f,
    Temperature = 0.95f
  };

  if (history != null) {
    request.History.AddRange(history);
  }

  return request;
}

Continue to rewrite the two chat methods by adding a history parameter.

public async Task<string> Chat(string prompt, List<Message>? history = null) {
  var resp = await _client.ChatAsync(GetRequest(prompt, history));
  return RenderText(resp.Response);
}

public async IAsyncEnumerable<string> StreamingChat(string prompt, List<Message>? history = null) {
  using var call = _client.StreamingChat(GetRequest(prompt, history));
  await foreach (var resp in call.ResponseStream.ReadAllAsync()) {
    yield return RenderText(resp.Response);
  }
}

Done.

Modifying the Python Code of the gRPC Server

First, let's see how ChatGLM2 handles the conversation input.

By debugging the official demo, I found that the history passed to the model is a list containing tuples, representing each conversation in a quirky format.

history = [('question1', 'answer1'), ('question2', 'answer2')]

However, AIHub's conversations follow the OpenAI style, with this format:

history = [
  {'role': 'user', 'content': 'question1'},
  {'role': 'assistant', 'content': 'answer1'},
  {'role': 'user', 'content': 'question2'},
  {'role': 'assistant', 'content': 'answer2'},
]

Now, we need to convert the OpenAI conversation format to the ChatGLM format.

Let me show you the code directly.

def messages_to_tuple_history(messages: List[chat_pb2.Message]):
    """Convert the chat message list to the nested list of tuples format required by ChatGLM."""
    history = []
    current_completion = ['', '']
    is_enter_completion = False

    
    for item in messages:
        if not is_enter_completion and item.role == 'user':
            is_enter_completion = True

        if is_enter_completion:
            if item.role == 'user':
                if len(current_completion[0]) > 0:
                    current_completion[0] = f"{current_completion[0]}\n\n{item.content}"
                else:
                    current_completion[0] = item.content
            if item.role == 'assistant':
                if len(current_completion[1]) > 0:
                    current_completion[1] = f"{current_completion[1]}\n\n{item.content}"
                else:
                    current_completion[1] = item.content

                is_enter_completion = False
                history.append((current_completion[0], current_completion[1]))
                current_completion = ['', '']

    return history

Currently, only the user and assistant roles are handled. Actually, OpenAI also has system and function roles. The system role is easier to handle and could be represented like this:

[('system prompt1', ''), ('system prompt2', '')]

However, I haven't tested it yet, and I don't need this for now, so I won't include it in the code.

Next, we continue to modify the two conversation methods.

class ChatService(chat_pb2_grpc.ChatHubServicer):
    def Chat(self, request: chat_pb2.ChatRequest, context):
        response, history = model.chat(
            tokenizer,
            request.prompt,
            history=messages_to_tuple_history(request.history),
            max_length=request.max_length,
            top_p=request.top_p,
            temperature=request.temperature)
        torch_gc()
        return chat_pb2.ChatReply(response=response)

    def StreamingChat(self, request: chat_pb2.ChatRequest, context):
        current_length = 0
        for response, history in model.stream_chat(
                tokenizer,
                request.prompt,
                history=messages_to_tuple_history(request.history),
                max_length=request.max_length,
                top_p=request.top_p,
                temperature=request.temperature,
                return_past_key_values=False):

            print(response[current_length:], end="", flush=True)
            yield chat_pb2.ChatReply(response=response)
            current_length = len(response)

        torch_gc()

By the way, remember to free up the GPU memory after each conversation is completed.

def torch_gc():
    if torch.cuda.is_available():
        with torch.cuda.device(CUDA_DEVICE):
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

That's it.

PS: You can use loguru for Python logging. It's very handy. I just discovered it recently.

Summary

Developing with gRPC-based invocation is still somewhat cumbersome, mainly because debugging can be tricky. I'm considering switching to a unified OpenAI interface approach for calling. Someone on GitHub has contributed an OpenAI-compatible interface for ChatGLM, which we can look into later.

However, for the visual aspect, we still need to continue using gRPC, as it offers better transmission efficiency. The reason large models can use HTTP's EventSource is that the data volume is relatively small, and a secondary reason is that the conversation is unidirectional—meaning users ask questions to the model, and the model does not proactively send information to the user.