Skip to content

Commit fd29b19

Browse files
committed
Properly support TokenStream return type in AiService
1 parent 4851856 commit fd29b19

File tree

2 files changed

+7
-0
lines changed

2 files changed

+7
-0
lines changed

core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/AiServicesProcessor.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,11 @@ public void handleDeclarativeServices(AiServicesRecorder recorder,
510510
// currently in one class either streaming or blocking model are supported, but not both
511511
// if we want to support it, the injectStreamingChatModelBean needs to be recorded per injection point
512512
for (MethodInfo method : declarativeAiServiceClassInfo.methods()) {
513+
if (LangChain4jDotNames.TOKEN_STREAM.equals(method.returnType().name())) {
514+
injectStreamingChatModelBean = true;
515+
continue;
516+
}
517+
513518
if (!DotNames.MULTI.equals(method.returnType().name())) {
514519
continue;
515520
}

core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/LangChain4jDotNames.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import dev.langchain4j.service.Moderate;
2323
import dev.langchain4j.service.Result;
2424
import dev.langchain4j.service.SystemMessage;
25+
import dev.langchain4j.service.TokenStream;
2526
import dev.langchain4j.service.UserMessage;
2627
import dev.langchain4j.service.UserName;
2728
import dev.langchain4j.web.search.WebSearchEngine;
@@ -45,6 +46,7 @@ public class LangChain4jDotNames {
4546
public static final DotName IMAGE_MODEL = DotName.createSimple(ImageModel.class);
4647
public static final DotName TOKEN_COUNT_ESTIMATOR = DotName.createSimple(TokenCountEstimator.class);
4748
public static final DotName CHAT_MESSAGE = DotName.createSimple(ChatMessage.class);
49+
public static final DotName TOKEN_STREAM = DotName.createSimple(TokenStream.class);
4850
public static final DotName OUTPUT_GUARDRAILS = DotName.createSimple(OutputGuardrails.class);
4951
public static final DotName INPUT_GUARDRAILS = DotName.createSimple(InputGuardrails.class);
5052
static final DotName AI_SERVICES = DotName.createSimple(AiServices.class);

0 commit comments

Comments
 (0)