Merge pull request #1005 from quarkiverse/token-stream

Properly support TokenStream return type in AiService
quarkiverse · Oct 24, 2024 · 7279fe5 · 7279fe5
2 parents c03cb43 + fd29b19
commit 7279fe5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/AiServicesProcessor.java b/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/AiServicesProcessor.java
@@ -530,6 +530,11 @@ public void handleDeclarativeServices(AiServicesRecorder recorder,
             // currently in one class either streaming or blocking model are supported, but not both
             // if we want to support it, the injectStreamingChatModelBean needs to be recorded per injection point
             for (MethodInfo method : declarativeAiServiceClassInfo.methods()) {
+                if (LangChain4jDotNames.TOKEN_STREAM.equals(method.returnType().name())) {
+                    injectStreamingChatModelBean = true;
+                    continue;
+                }
+
                 if (!DotNames.MULTI.equals(method.returnType().name())) {
                     continue;
                 }

diff --git a/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/LangChain4jDotNames.java b/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/LangChain4jDotNames.java
@@ -22,6 +22,7 @@
 import dev.langchain4j.service.Moderate;
 import dev.langchain4j.service.Result;
 import dev.langchain4j.service.SystemMessage;
+import dev.langchain4j.service.TokenStream;
 import dev.langchain4j.service.UserMessage;
 import dev.langchain4j.service.UserName;
 import dev.langchain4j.web.search.WebSearchEngine;
@@ -45,6 +46,7 @@ public class LangChain4jDotNames {
     public static final DotName IMAGE_MODEL = DotName.createSimple(ImageModel.class);
     public static final DotName TOKEN_COUNT_ESTIMATOR = DotName.createSimple(TokenCountEstimator.class);
     public static final DotName CHAT_MESSAGE = DotName.createSimple(ChatMessage.class);
+    public static final DotName TOKEN_STREAM = DotName.createSimple(TokenStream.class);
     public static final DotName OUTPUT_GUARDRAILS = DotName.createSimple(OutputGuardrails.class);
     public static final DotName INPUT_GUARDRAILS = DotName.createSimple(InputGuardrails.class);
     static final DotName AI_SERVICES = DotName.createSimple(AiServices.class);