ops-api/src/main/java/com/qmth/ops/biz/ai/client/DocApiClient.java

package com.qmth.ops.biz.ai.client;

import com.qmth.boot.core.ai.model.AiConstants;
import com.qmth.boot.core.ai.model.doc.ParseDocTask;
import com.qmth.boot.core.ai.model.doc.ParseDocTaskResult;
import com.qmth.boot.core.exception.StatusException;
import com.qmth.boot.core.rateLimit.service.RateLimiter;
import com.qmth.boot.core.rateLimit.service.impl.MemoryRateLimiter;
import com.qmth.boot.tools.codec.CodecUtils;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.time.Duration;

/**
 * 文档解析类接口基础实现
 */
public abstract class DocApiClient {

    private static final Logger log = LoggerFactory.getLogger(DocApiClient.class);

    private final DocApiConfig config;

    private final OkHttpClient client;

    private RateLimiter queryRateLimiter;

    public DocApiClient(DocApiConfig config) {
        this.config = config;

        OkHttpClient.Builder builder = new OkHttpClient.Builder().connectionPool(new ConnectionPool())
                .connectTimeout(Duration.ofSeconds(10)).readTimeout(Duration.ofSeconds(50));
        Interceptor interceptor = getInterceptor();
        if (interceptor != null) {
            builder.addInterceptor(interceptor);
        }
        this.client = builder.build();

        if (config.getQps() > 0) {
            this.queryRateLimiter = new MemoryRateLimiter(config.getQps(), 1000);
        }
    }

    public abstract ParseDocTask parseDocTask(byte[] fileData, String fileName) throws Exception;

    public abstract ParseDocTaskResult parseDocTaskQuery(String taskId) throws Exception;

    protected DocApiConfig getConfig() {
        return config;
    }

    protected OkHttpClient getClient() {
        return client;
    }

    protected RateLimiter getQueryRateLimiter() {
        return queryRateLimiter;
    }

    protected Interceptor getInterceptor() {
        return null;
    }

    public String encodeTaskId(String taskId) {
        String str = getConfig().getSupplier() + AiConstants.PARAM_SPLIT + taskId;
        return CodecUtils.toBase64(str.getBytes(StandardCharsets.UTF_8));
    }

    public String[] decodeTaskId(String encodeStr) {
        String str = new String(CodecUtils.fromBase64(encodeStr), StandardCharsets.UTF_8);
        String[] values = str.split(AiConstants.PARAM_SPLIT);
        if (values.length != 3) {
            throw new StatusException("taskId值无效！");
        }
        return values;
    }

    protected byte[] download(String url) {
        Request request = new Request.Builder().url(url).get().build();
        try (Response response = this.getClient().newCall(request).execute();) {
            ResponseBody respBody = response.body();
            if (response.isSuccessful() && respBody != null) {
                return respBody.bytes();
            }
            log.error("获取文件内容失败！responseCode:{}", response.code());
            throw new StatusException("获取文件内容失败！");
        } catch (Exception e) {
            throw new StatusException("获取文件内容失败！", e);
        }
    }

}