Langchain4j工具调用获取不到ThreadLocal

问题描述

最近用SpringBoot开发个人项目使用了Langchain4j 1.12.2，开发一个功能，让AI输出当前用户信息，用户信息在请求拦截器中存入了ThreadLocal，工具调用时我希望拿到用户信息交给AI，但是死活拿不到，于是我给langchain4j官方github提了issue，最终找到

解决方案

① 调用ai时将数据封装到InvocationParameters中

② 工具方法上加上 InvocationParameters 参数

这样即可在运行时获取到上游传来的参数

使用案例：

java 复制代码

import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.invocation.InvocationContext;
import dev.langchain4j.invocation.InvocationParameters;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import org.junit.jupiter.api.Test;

import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.concurrent.CountDownLatch;

public class Langchain4jParameterTest {

    public interface MyAiService {
        TokenStream chat(@UserMessage String userMessage, InvocationParameters parameters);
    }

    @Tool( name = "getNowTime", value = "get now time")
    public String getNowTime(InvocationParameters parameters) {
        String user = parameters.get("user");
        System.out.println("\n工具调用：用户"+ user +"获取当前时间");
        return LocalDateTime.now().toString();
    }

    @Test
    public void test() {
        OpenAiStreamingChatModel streamingChatModel = OpenAiStreamingChatModel.builder()
                .modelName("模型名称")
                .apiKey("填上你的apiKey")
                .baseUrl("访问的baseUrl")
                .build();
        MyAiService aiService = AiServices.builder(MyAiService.class)
                .streamingChatModel(streamingChatModel)
                .tools(this)
                .build();
        HashMap<String, Object> map = new HashMap<>();
        map.put("user", "大大帅");
        InvocationParameters invocationParameters = InvocationParameters.from(map);
        TokenStream tokenStream = aiService.chat("现在几点了", invocationParameters);
        CountDownLatch countDownLatch = new CountDownLatch(1);
        tokenStream.onPartialResponse(chatResponse -> {
            System.out.print(chatResponse);
        });
        tokenStream.onCompleteResponse(chatResponse -> {
            countDownLatch.countDown();
        });
        tokenStream.onError(throwable -> {
            System.out.println("【错误: " + throwable.getMessage() + "】");
            countDownLatch.countDown();
        });
        tokenStream.start();
        try {
            countDownLatch.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
}

输出结果：

可以看到已经获取到传递的user了

官方参考文档

https://docs.langchain4j.dev/tutorials/tools#invocationparameters

思考总结

在我debug langchain4j的源码后发现底层请求使用了JDK（JUC包下）的ForkJoinPool的线程池，导致响应后进行工具回调不在是原先的线程，因此获取不到ThreadLocal，官方提供的设置自定义线程池的api 在创建AiService时通过executeToolsConcurrently(executor)方法也无法实现ThreadLocal的正确传递，本质上只是在高层替换了线程池，在底层使用还是jdk的ForkJoinPool，简言之是当前线程 -> 调用自定义executor -> 调用ForkJoinPool -> 工具调用。

有些第三方AIChatModel可以传递一次InheribleThreadLocal的值，这仅仅是因为线程池新创建线程复制了一份父线程的InheribleThreadLocal值，后续线程池复用将导致无法获取到正确的InheribleThreadLocal，如：

① QwenStreamingChatModel能获取到第一次的InheribleThreadLocal

② ZhipuAiStreamingChatModel获取不到InheribleThreadLocal

③ OpenAiStreamingChatModel也获取不到InheribleThreadLocal

即使使用阿里的ttl对线程池进行包装+TransmittableThreadLocal也无法实现线程传值，我目前找到最有效的方案就是通过InvocationParameters向工具调用时传递数据，当然也可以通过InvocationContext获取（内部有一个invocationParameters属性）