SpringでOpenAIのストリーミングアウトプットに対応したアプリケーションを作成する

SpringのアプリケーションでOpenAIのストリーミングを有効にしたコードを紹介します。 なお、この記事を書いている時点では、ハック的なコードが多いため、あくまでこのブログの執筆時点での方法と考えてくだい。

まずはSpring AIのお話

OpenAIのAPIを利用したJava/Springアプリケーション開発を支援してくれるのがSpring AIプロジェクト.

Getting Startedに従うとこんな感じで、生成AIが行えるシンプルなAPIエンドポイントを作ることができます。 以下は生成AI(私はLLama2を利用)に対して “Tell me a joke"と聞いた結果

 1% curl "localhost:8080/ai/simple?message=Tell%20me%20a%20joke"
 2  Sure, here's one for you:
 3
 4Why couldn't the bicycle stand up by itself?
 5
 6Because it was two-tired!
 7
 8Get it? Two-tired... like a bicycle tire... ha ha ha!
 9
10Hope that brought a smile to your face!

解決したい課題

Spring AIはこの執筆時点では、ストリーム出力をサポートしていません。以下のIssueで議論されています。

https://github.com/spring-projects/spring-ai/issues/116

結果として、生成AIの出力が全部終わるまでユーザーは待たなくてはいけません。 たとえば、さきほどのプロンプトを"Tell me a very long joke"と聞いてやりなおすと、以下のような結果です。

 1% time curl "localhost:8080/ai/simple?message=Tell%20me%20a%20very%20long%20joke"
 2  Sure! Here's a very long joke for you:
 3
 4One day, a man walked into a library and asked the librarian, "Do you have any books on Pavlov's dogs and Schrödinger's cat?"
 5
 6The librarian replied, "It rings a bell, but I'm not sure if it's here or not."
 7
 8The man thought for a moment and said, "Well, I'm looking for a book that explores the concept of conditioned response and the idea that a cat can be both alive and dead at the same time."
 9
10The librarian scratched her head and said, "Let me see... I think I have just the book for you. It's called 'Pavlov's Cats and Schrödinger's Dogs.'"
11
12The man's eyes widened in surprise and he exclaimed, "That's exactly what I'm looking for! But can I find it on the shelf?"
13
14The librarian smiled and said, "It's a bit of a puzzle, but if you look carefully, you should be able to find it. Just remember, the book is neither here nor there, but it's definitely somewhere in the library."
15
16The man thought for a moment and then asked, "But how do I know if I've found the right book if it's both here and not here at the same time?"
17
18The librarian chuckled and said, "Well, that's the million-dollar question. But I think you'll know it when you see it. The book will be neither on the shelf nor not on the shelf, but somewhere in between. And when you open it, you'll see that it's both a book about Pavlov's dogs and Schrödinger's cat, and yet, it's not either of those things at the same time. It's a bit of a paradox, but that's the beauty of it."
19
20The man nodded thoughtfully and began to search the shelves, determined to find the elusive book. After a few minutes of searching, he finally found it nestled between two other books that were neither here nor there. He opened it up and was amazed to find that it was both a book about Pavlov's dogs and Schrödinger's cat, and yet, it was neither of those things at the same time.
21
22The end.
23
24I hope that long joke brought a smile to your face!
25curl "localhost:8080/ai/simple?message=Tell%20me%20a%20very%20long%20joke"  0.00s user 0.01s system 0% cpu 13.769 total

cpu 13.769 totalにあるよう、この結果には13秒近く待たされたことになります。 これをWebアプリケーション化しても、ユーザーからすると、10秒以上無応答になってしまい"壊れた?“を疑われてしまいます。

本来この課題を解決するのがストリーム出力です。生成AIが完全に出力するのをまつのではなく、現在生成した箇所を表示することによってレスポンスが即座にユーザーにかえってきていることを見せるためのものです。

SpringAIは、今のところはまだ、ストリーム出力に対応していませんが、ハック的にストリーム出力に対応したアプリケーションを作ってみたいと思います。

コード

https://github.com/mhoshi-vm/StreamSpringExample

手順

以後実装の手順を書いていきます。

下準備

まずは、Getting Startedに従った通常のアプリケーションを作ります。

ちょっとしたハック

これは、私の環境依存なのですが私の環境のAPIエンドポイントが https://xxxx/api がプリフィックスになっているのですが、これが Spring AIで使われているクライアントがうまく扱えない状況です。 ここのIssueに記載しています。

https://github.com/TheoKanning/openai-java/issues/370

この事象に対応するために、以下のコードでAIエンドポイントを上がいています。

https://github.com/mhoshi-vm/StreamSpringExample/blob/main/src/main/java/com/example/streamspringexample/localOpenAi/MyOpenAiApi.java

本家のOpenAIを使う場合は不要な手順ですので、ご注意ください。

Starter Web と Starter WebFlux の両方の依存関係を定義

pom.xml に以下の依存関係を追加していきます。

1        <dependency>
2            <groupId>org.springframework.boot</groupId>
3            <artifactId>spring-boot-starter-web</artifactId>
4        </dependency>
5        <dependency>
6            <groupId>org.springframework.boot</groupId>
7            <artifactId>spring-boot-starter-webflux</artifactId>
8        </dependency>

なお、本当は spring-boot-starter-webflux 不要なはずです。細かい検証ができていなく、原因が不明なのですが、執筆時点では、この対応が必要だった点を追記しておきます。

OpenAI Client の拡張および新フアンクション定義

org.springframework.ai.openai.client.OpenAiClient を拡張して、ストリームに対応したメソッドを定義しています。

 1public class MyOpenAiClient extends OpenAiClient {
 2
 3    private final OpenAiService openAiService;
 4
 5    public MyOpenAiClient(OpenAiService openAiService) {
 6        super(openAiService);
 7        this.openAiService = openAiService;
 8    }
 9
10    public Flowable<String> generateStream(String prompt) {
11
12        ChatCompletionRequest completionRequest = ChatCompletionRequest.builder()
13                .model(super.getModel())
14                .temperature(super.getTemperature())
15                .messages(List.of(new ChatMessage("user", prompt)))
16                .build();
17
18        return openAiService.streamChatCompletion(completionRequest)
19                .filter(completionChunk ->
20                        completionChunk.getChoices().get(0) != null && completionChunk.getChoices().get(0).getMessage() != null && completionChunk.getChoices().get(0).getMessage().getContent() != null).map(
21                        completionChunk -> completionChunk.getChoices().get(0).getMessage().getContent());
22    }
23
24}

最終的なコードはこちら。

https://github.com/mhoshi-vm/StreamSpringExample/blob/main/src/main/java/com/example/streamspringexample/localOpenAi/MyOpenAiClient.java

Controller の定義

以下のようなコントローラーを定義します。

 1@RestController
 2public class SimpleController {
 3
 4    private final MyOpenAiClient aiClient;
 5
 6    public SimpleController(MyOpenAiClient aiClient) {
 7        this.aiClient = aiClient;
 8    }
 9
10    @GetMapping(path = "/ai/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
11    public Flowable<String> streamCompletion(@RequestParam(value = "message", defaultValue = "Tell me a very long joke") String message) {
12        return aiClient.generateStream(message);
13    }
14}

ポイントとして、MediaType を “TEXT_EVENT_STREAM_VALUE” として定義していることを返り値を Flowable<Stream> にすることです。 こうすることで、ストリーム型のアウトプットを使うことができます。

最終的なコードはこちら。

https://github.com/mhoshi-vm/StreamSpringExample/blob/main/src/main/java/com/example/streamspringexample/controller/SimpleController.java

実装としては以上です。

実行例

こんな感じで実行すると、10秒待つことなく、いきなりアウトプットが得られます。

 1machih@machihXCV5C StreamSpringExample % time curl "localhost:8080/ai/stream?message=Tell%20me%20a%20very%20long%20joke" -H 'accept: text/event-stream'
 2data: 
 3
 4data: Sure
 5
 6data:,
 7
 8data: here
 9
10data:'
11
12data:s
13
14data: a
15
16data: very
17
18data: long
19
20data: jo
21
22data:ke
23
24data: for
25
26...

おまけ

コードを改良して、localhost:8080 でブラウザでアクセスすると、ヌルヌルと出力がみれるようにしました。