001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.model.chat;
017
018import com.agentsflex.core.message.AiMessage;
019import com.agentsflex.core.model.chat.response.AiMessageResponse;
020import com.agentsflex.core.model.client.StreamContext;
021import com.agentsflex.core.observability.Observability;
022import io.opentelemetry.api.common.AttributeKey;
023import io.opentelemetry.api.common.Attributes;
024import io.opentelemetry.api.metrics.DoubleHistogram;
025import io.opentelemetry.api.metrics.LongCounter;
026import io.opentelemetry.api.metrics.Meter;
027import io.opentelemetry.api.trace.Span;
028import io.opentelemetry.api.trace.StatusCode;
029import io.opentelemetry.api.trace.Tracer;
030import io.opentelemetry.context.Scope;
031
032import java.util.concurrent.atomic.AtomicBoolean;
033
034public class ChatObservabilityInterceptor implements ChatInterceptor {
035
036    private static final Tracer TRACER = Observability.getTracer();
037    private static final Meter METER = Observability.getMeter();
038
039    private static final LongCounter LLM_REQUEST_COUNT = METER.counterBuilder("llm.request.count")
040        .setDescription("Total number of LLM requests")
041        .build();
042
043    private static final DoubleHistogram LLM_LATENCY_HISTOGRAM = METER.histogramBuilder("llm.request.latency")
044        .setDescription("LLM request latency in seconds")
045        .setUnit("s")
046        .build();
047
048    private static final LongCounter LLM_ERROR_COUNT = METER.counterBuilder("llm.request.error.count")
049        .setDescription("Total number of LLM request errors")
050        .build();
051
052    private static final int MAX_RESPONSE_LENGTH_FOR_SPAN = 500;
053
054    @Override
055    public AiMessageResponse intercept(BaseChatModel<?> chatModel, ChatContext context, SyncChain chain) {
056
057        ChatConfig config = chatModel.getConfig();
058
059        if (config == null || !config.isObservabilityEnabled() || !Observability.isEnabled()) {
060            return chain.proceed(chatModel, context);
061        }
062
063        String provider = config.getProvider();
064        String model = config.getModel();
065        String operation = "chat";
066
067        Span span = TRACER.spanBuilder(provider + "." + operation)
068            .setAttribute("llm.provider", provider)
069            .setAttribute("llm.model", model)
070            .setAttribute("llm.operation", operation)
071            .startSpan();
072
073
074        long startTimeNanos = System.nanoTime();
075
076        try (Scope ignored = span.makeCurrent()) {
077
078            AiMessageResponse response = chain.proceed(chatModel, context);
079            boolean success = (response != null) && !response.isError();
080
081            if (success) {
082                enrichSpan(span, response.getMessage());
083            }
084
085            recordMetrics(provider, model, operation, success, startTimeNanos);
086            return response;
087
088        } catch (Exception e) {
089            recordError(span, e, provider, model, operation, startTimeNanos);
090            throw e;
091        } finally {
092            span.end();
093        }
094    }
095
096    @Override
097    public void interceptStream(
098        BaseChatModel<?> chatModel,
099        ChatContext context,
100        StreamResponseListener originalListener,
101        StreamChain chain) {
102
103        ChatConfig config = chatModel.getConfig();
104
105        if (config == null || !config.isObservabilityEnabled()) {
106            chain.proceed(chatModel, context, originalListener);
107            return;
108        }
109
110        String provider = config.getProvider();
111        String model = config.getModel();
112        String operation = "chatStream";
113
114        Span span = TRACER.spanBuilder(provider + "." + operation)
115            .setAttribute("llm.provider", provider)
116            .setAttribute("llm.model", model)
117            .setAttribute("llm.operation", operation)
118            .startSpan();
119
120        Scope scope = span.makeCurrent();
121        long startTimeNanos = System.nanoTime();
122
123        AtomicBoolean recorded = new AtomicBoolean(false);
124
125        StreamResponseListener wrappedListener = new StreamResponseListener() {
126            @Override
127            public void onStart(StreamContext context) {
128                originalListener.onStart(context);
129            }
130
131            @Override
132            public void onMessage(StreamContext context, AiMessageResponse response) {
133                originalListener.onMessage(context, response);
134            }
135
136
137            @Override
138            public void onFailure(StreamContext context, Throwable throwable) {
139                safeRecord(false, throwable);
140                originalListener.onFailure(context, throwable);
141            }
142
143            @Override
144            public void onStop(StreamContext context) {
145                boolean success = !context.isError();
146                if (success) {
147                    enrichSpan(span, context.getFullMessage());
148                }
149                safeRecord(success, null);
150                originalListener.onStop(context);
151            }
152
153            private void safeRecord(boolean success, Throwable throwable) {
154                if (recorded.compareAndSet(false, true)) {
155                    if (throwable != null) {
156                        span.setStatus(StatusCode.ERROR, throwable.getMessage());
157                        span.recordException(throwable);
158                    }
159                    span.end();
160                    scope.close();
161                    recordMetrics(provider, model, operation, success, startTimeNanos);
162                }
163            }
164        };
165
166        try {
167            chain.proceed(chatModel, context, wrappedListener);
168        } catch (Exception e) {
169            if (recorded.compareAndSet(false, true)) {
170                recordError(span, e, provider, model, operation, startTimeNanos);
171            }
172            scope.close();
173            throw e;
174        }
175    }
176
177    private void enrichSpan(Span span, AiMessage msg) {
178        if (msg != null) {
179            span.setAttribute("llm.total_tokens", msg.getEffectiveTotalTokens());
180            String content = msg.getContent();
181            if (content != null) {
182                span.setAttribute("llm.response",
183                    content.substring(0, Math.min(content.length(), MAX_RESPONSE_LENGTH_FOR_SPAN)));
184            }
185        }
186    }
187
188    private void recordMetrics(String provider, String model, String operation, boolean success, long startTimeNanos) {
189        double latencySeconds = (System.nanoTime() - startTimeNanos) / 1_000_000_000.0;
190        Attributes attrs = Attributes.of(
191            AttributeKey.stringKey("llm.provider"), provider,
192            AttributeKey.stringKey("llm.model"), model,
193            AttributeKey.stringKey("llm.operation"), operation,
194            AttributeKey.stringKey("llm.success"), String.valueOf(success)
195        );
196        LLM_REQUEST_COUNT.add(1, attrs);
197        LLM_LATENCY_HISTOGRAM.record(latencySeconds, attrs);
198        if (!success) {
199            LLM_ERROR_COUNT.add(1, attrs);
200        }
201    }
202
203    private void recordError(Span span, Exception e, String provider, String model, String operation, long startTimeNanos) {
204        span.setStatus(StatusCode.ERROR, e.getMessage());
205        span.recordException(e);
206        span.end();
207        // Scope 会在 finally 或 safeRecord 中关闭
208        recordMetrics(provider, model, operation, false, startTimeNanos);
209    }
210}