Make response handling more generic. Resulting in generic cost handling.

Author: alesj

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

@@ -2,6 +2,12 @@
 
 import java.util.Optional;
 
+import jakarta.inject.Singleton;
+
+import org.jboss.jandex.DotName;
+
+import io.quarkiverse.langchain4j.cost.CostEstimatorResponseListener;
+import io.quarkiverse.langchain4j.deployment.config.LangChain4jBuildConfig;
 import io.quarkiverse.langchain4j.runtime.listeners.MetricsChatModelListener;
 import io.quarkiverse.langchain4j.runtime.listeners.SpanChatModelListener;
 import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
@@ -14,6 +20,20 @@
 
 public class ListenersProcessor {
 
+    @BuildStep
+    public void costListener(
+            LangChain4jBuildConfig config,
+            BuildProducer<AdditionalBeanBuildItem> additionalBeanProducer) {
+        if (config.costListener()) {
+            additionalBeanProducer.produce(
+                    AdditionalBeanBuildItem.builder()
+                            .addBeanClass(CostEstimatorResponseListener.class)
+                            .setDefaultScope(DotName.createSimple(Singleton.class))
+                            .setUnremovable()
+                            .build());
+        }
+    }
+
     @BuildStep
     public void spanListeners(Capabilities capabilities,
             Optional<MetricsCapabilityBuildItem> metricsCapability,

core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/config/LangChain4jBuildConfig.java

@@ -43,6 +43,12 @@ public interface LangChain4jBuildConfig {
     @WithDefault("true")
     boolean responseSchema();
 
+    /**
+     * Configuration property to enable or disable generic cost listener
+     */
+    @WithDefault("false")
+    boolean costListener();
+
     interface BaseConfig {
         /**
          * Chat model

core/runtime/src/main/java/io/quarkiverse/langchain4j/cost/CostEstimatorResponseListener.java

@@ -0,0 +1,55 @@
+package io.quarkiverse.langchain4j.cost;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+
+import jakarta.inject.Inject;
+
+import dev.langchain4j.model.output.TokenUsage;
+import io.quarkiverse.langchain4j.response.ResponseListener;
+import io.quarkiverse.langchain4j.response.ResponseRecord;
+import io.quarkus.arc.All;
+import io.smallrye.common.annotation.Experimental;
+
+/**
+ * Allows for user code to provide a custom strategy for estimating the cost of API calls
+ */
+@Experimental("This feature is experimental and the API is subject to change")
+public class CostEstimatorResponseListener implements ResponseListener {
+
+    private final CostEstimatorService service;
+    private final List<CostListener> listeners;
+
+    @Inject
+    public CostEstimatorResponseListener(CostEstimatorService service, @All List<CostListener> listeners) {
+        this.service = service;
+        this.listeners = new ArrayList<>(listeners);
+        this.listeners.sort(Comparator.comparingInt(CostListener::order));
+    }
+
+    @Override
+    public void onResponse(ResponseRecord rr) {
+        String model = rr.model();
+        TokenUsage tokenUsage = rr.tokenUsage();
+        CostEstimator.CostContext context = new MyCostContext(tokenUsage, model);
+        Cost cost = service.estimate(context);
+        if (cost != null) {
+            for (CostListener cl : listeners) {
+                cl.handleCost(model, tokenUsage, cost);
+            }
+        }
+    }
+
+    private record MyCostContext(TokenUsage tokenUsage, String model) implements CostEstimator.CostContext {
+        @Override
+        public Integer inputTokens() {
+            return tokenUsage().inputTokenCount();
+        }
+
+        @Override
+        public Integer outputTokens() {
+            return tokenUsage().outputTokenCount();
+        }
+    }
+}

core/runtime/src/main/java/io/quarkiverse/langchain4j/cost/CostEstimatorService.java

@@ -28,10 +28,13 @@ public CostEstimatorService(@All List<CostEstimator> costEstimators) {
     public Cost estimate(ChatModelResponseContext response) {
         TokenUsage tokenUsage = response.response().tokenUsage();
         CostEstimator.CostContext costContext = new MyCostContext(tokenUsage, response);
+        return estimate(costContext);
+    }
 
+    public Cost estimate(CostEstimator.CostContext context) {
         for (CostEstimator costEstimator : costEstimators) {
-            if (costEstimator.supports(costContext)) {
-                CostEstimator.CostResult costResult = costEstimator.estimate(costContext);
+            if (costEstimator.supports(context)) {
+                CostEstimator.CostResult costResult = costEstimator.estimate(context);
                 if (costResult != null) {
                     BigDecimal totalCost = costResult.inputTokensCost().add(costResult.outputTokensCost());
                     return new Cost(totalCost, costResult.currency());

core/runtime/src/main/java/io/quarkiverse/langchain4j/cost/CostListener.java

@@ -0,0 +1,14 @@
+package io.quarkiverse.langchain4j.cost;
+
+import dev.langchain4j.model.output.TokenUsage;
+
+/**
+ * Allows for user code to handle estimate cost; e.g. some simple accounting
+ */
+public interface CostListener {
+    void handleCost(String model, TokenUsage tokenUsage, Cost cost);
+
+    default int order() {
+        return 0;
+    }
+}

core/runtime/src/main/java/io/quarkiverse/langchain4j/response/ResponseInterceptor.java

@@ -0,0 +1,47 @@
+package io.quarkiverse.langchain4j.response;
+
+import java.util.Map;
+
+import jakarta.annotation.Priority;
+import jakarta.interceptor.AroundInvoke;
+import jakarta.interceptor.Interceptor;
+import jakarta.interceptor.InvocationContext;
+
+import dev.langchain4j.data.message.AiMessage;
+import dev.langchain4j.model.chat.listener.ChatModelResponse;
+import dev.langchain4j.model.chat.response.ChatResponse;
+import dev.langchain4j.model.output.Response;
+
+/**
+ * Simple (Chat)Response interceptor, to be applied directly on the model.
+ */
+@Interceptor
+@ResponseInterceptorBinding
+@Priority(0)
+public class ResponseInterceptor extends ResponseInterceptorBase {
+
+    @AroundInvoke
+    public Object intercept(InvocationContext context) throws Exception {
+        Object result = context.proceed();
+        ResponseRecord rr = null;
+        if (result instanceof Response<?> response) {
+            Object content = response.content();
+            if (content instanceof AiMessage am) {
+                rr = new ResponseRecord(getModel(context.getTarget()), am, response.tokenUsage(), response.finishReason(),
+                        response.metadata());
+            }
+        } else if (result instanceof ChatResponse response) {
+            rr = new ResponseRecord(getModel(context.getTarget()), response.aiMessage(), response.tokenUsage(),
+                    response.finishReason(), Map.of());
+        } else if (result instanceof ChatModelResponse response) {
+            rr = new ResponseRecord(response.model(), response.aiMessage(), response.tokenUsage(), response.finishReason(),
+                    Map.of("id", response.id()));
+        }
+        if (rr != null) {
+            for (ResponseListener l : getListeners()) {
+                l.onResponse(rr);
+            }
+        }
+        return result;
+    }
+}

core/runtime/src/main/java/io/quarkiverse/langchain4j/response/ResponseInterceptorBase.java

@@ -0,0 +1,41 @@
+package io.quarkiverse.langchain4j.response;
+
+import java.lang.reflect.Method;
+import java.util.Comparator;
+import java.util.List;
+
+import jakarta.enterprise.inject.Any;
+import jakarta.enterprise.inject.spi.CDI;
+
+/**
+ * Simple (Chat)Response interceptor base, to be applied directly on the model.
+ */
+public abstract class ResponseInterceptorBase {
+
+    private volatile String model;
+    private volatile List<ResponseListener> listeners;
+
+    // TODO -- uh uh ... reflection ... puke
+    protected String getModel(Object target) {
+        if (model == null) {
+            try {
+                Class<?> clazz = target.getClass();
+                Method method = clazz.getMethod("modelName");
+                model = (String) method.invoke(target);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        }
+        return model;
+    }
+
+    protected List<ResponseListener> getListeners() {
+        if (listeners == null) {
+            listeners = CDI.current().select(ResponseListener.class, Any.Literal.INSTANCE)
+                    .stream()
+                    .sorted(Comparator.comparing(ResponseListener::order))
+                    .toList();
+        }
+        return listeners;
+    }
+}

core/runtime/src/main/java/io/quarkiverse/langchain4j/response/ResponseInterceptorBinding.java

@@ -0,0 +1,14 @@
+package io.quarkiverse.langchain4j.response;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import jakarta.interceptor.InterceptorBinding;
+
+@InterceptorBinding
+@Target({ ElementType.TYPE, ElementType.METHOD })
+@Retention(RetentionPolicy.RUNTIME)
+public @interface ResponseInterceptorBinding {
+}

core/runtime/src/main/java/io/quarkiverse/langchain4j/response/ResponseInterceptorBindingSource.java

@@ -0,0 +1,5 @@
+package io.quarkiverse.langchain4j.response;
+
+@ResponseInterceptorBinding
+public abstract class ResponseInterceptorBindingSource {
+}

core/runtime/src/main/java/io/quarkiverse/langchain4j/response/ResponseListener.java

@@ -0,0 +1,12 @@
+package io.quarkiverse.langchain4j.response;
+
+/**
+ * Simple ResponseRecord listener, to be implemented by the (advanced) users.
+ */
+public interface ResponseListener {
+    void onResponse(ResponseRecord response);
+
+    default int order() {
+        return 0;
+    }
+}