# Agent数据分析V3.0.1.7

Junjie

2 天以前 63b01db83d9aad8a15276b4236a9a22e4aeef065

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
package com.zy.ai.service.impl;
 
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.zy.ai.entity.AiTokenUsage;
import com.zy.ai.mapper.AiTokenUsageMapper;
import com.zy.ai.service.AiTokenUsageService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
 
@Slf4j
@Service("aiTokenUsageService")
public class AiTokenUsageServiceImpl extends ServiceImpl<AiTokenUsageMapper, AiTokenUsage>
        implements AiTokenUsageService {
 
    @Override
    public void incrementTokens(long promptTokens, long completionTokens, long totalTokens, long callCount) {
        if (promptTokens <= 0 && completionTokens <= 0 && totalTokens <= 0 && callCount <= 0) {
            return;
        }
        try {
            int rows = baseMapper.incrementTokens(promptTokens, completionTokens, totalTokens, callCount);
            if (rows == 0) {
                // Row doesn't exist, create it
                AiTokenUsage usage = new AiTokenUsage();
                usage.setId(1);
                usage.setPromptTokens(promptTokens);
                usage.setCompletionTokens(completionTokens);
                usage.setTotalTokens(totalTokens);
                usage.setLlmCallCount(callCount);
                save(usage);
            }
        } catch (Exception e) {
            log.warn("Failed to increment AI token usage: {}", e.getMessage());
        }
    }
}