content: /*
   *  Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com).
   *  <p>
   *  Licensed under the Apache License, Version 2.0 (the "License");
   *  you may not use this file except in compliance with the License.
   *  You may obtain a copy of the License at
   *  <p>
   *  http://www.apache.org/licenses/LICENSE-2.0
   *  <p>
   *  Unless required by applicable law or agreed to in writing, software
   *  distributed under the License is distributed on an "AS IS" BASIS,
   *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   *  See the License for the specific language governing permissions and
   *  limitations under the License.
   */
  package com.agentsflex.llm.spark;
  
  import com.agentsflex.core.document.Document;
  import com.agentsflex.core.llm.BaseLlm;
  import com.agentsflex.core.llm.ChatContext;
  import com.agentsflex.core.llm.ChatOptions;
  import com.agentsflex.core.llm.StreamResponseListener;
  import com.agentsflex.core.llm.client.BaseLlmClientListener;
  import com.agentsflex.core.llm.client.HttpClient;
  import com.agentsflex.core.llm.client.LlmClient;
  import com.agentsflex.core.llm.client.LlmClientListener;
  import com.agentsflex.core.llm.client.impl.WebSocketClient;
  import com.agentsflex.core.llm.embedding.EmbeddingOptions;
  import com.agentsflex.core.llm.response.AbstractBaseMessageResponse;
  import com.agentsflex.core.llm.response.AiMessageResponse;
  import com.agentsflex.core.message.AiMessage;
  import com.agentsflex.core.parser.AiMessageParser;
  import com.agentsflex.core.prompt.Prompt;
  import com.agentsflex.core.store.VectorData;
  import com.agentsflex.core.util.SleepUtil;
  import com.agentsflex.core.util.StringUtil;
  import com.alibaba.fastjson.JSONPath;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.nio.ByteBuffer;
  import java.nio.ByteOrder;
  import java.util.Base64;
  import java.util.concurrent.CountDownLatch;
  
  public class SparkLlm extends BaseLlm<SparkLlmConfig> {
  
      private static final Logger logger = LoggerFactory.getLogger(SparkLlm.class);
      public AiMessageParser aiMessageParser = SparkLlmUtil.getAiMessageParser();
  
      private final HttpClient httpClient = new HttpClient();
  
  
      public SparkLlm(SparkLlmConfig config) {
          super(config);
      }
  
      @Override
      public VectorData embed(Document document, EmbeddingOptions options) {
          return embed(document, options, 0);
      }
  
  
      public VectorData embed(Document document, EmbeddingOptions options, int tryTimes) {
          String payload = SparkLlmUtil.embedPayload(config, document);
          String resp = httpClient.post(SparkLlmUtil.createEmbedURL(config), null, payload);
          if (StringUtil.noText(resp)) {
              logger.error("Could not get embed data" + document);
              return null;
          }
  
          Integer code = JSONPath.read(resp, "$.header.code", Integer.class);
          if (code == null) {
              logger.error(resp);
              return null;
          }
  
          if (code != 0) {
              //11202	授权错误：秒级流控超限。秒级并发超过授权路数限制
              if (code.equals(11202) && tryTimes < 3) {
                  SleepUtil.sleep(config.getConcurrencyLimitSleepMillis());
                  return embed(document, options, tryTimes + 1);
              } else {
                  logger.error(resp);
                  return null;
              }
          }
  
          String text = JSONPath.read(resp, "$.payload.feature.text", String.class);
          if (StringUtil.noText(text)) {
              logger.error(resp);
              return null;
          }
  
          byte[] buffer = Base64.getDecoder().decode(text);
          double[] vector = new double[buffer.length / 4];
          for (int i = 0; i < vector.length; i++) {
              int n = i * 4;
              vector[i] = ByteBuffer.wrap(buffer, n, 4).order(ByteOrder.LITTLE_ENDIAN).getFloat();
          }
  
          VectorData vectorData = new VectorData();
          vectorData.setVector(vector);
          return vectorData;
      }
  
  
      @Override
      public AiMessageResponse chat(Prompt prompt, ChatOptions options) {
          CountDownLatch latch = new CountDownLatch(1);
          Throwable[] failureThrowable = new Throwable[1];
          AiMessageResponse[] messageResponse = {null};
  
          waitResponse(prompt, options, messageResponse, latch, failureThrowable);
  
          AiMessageResponse response = messageResponse[0];
          Throwable fialureThrowable = failureThrowable[0];
  
          if (response == null) {
              if (fialureThrowable != null) {
                  response = new AiMessageResponse(prompt, "", null);
              } else {
                  return null;
              }
          }
  
          if (fialureThrowable != null || response.getMessage() == null) {
              response.setError(true);
              if (fialureThrowable != null) {
                  response.setErrorMessage(fialureThrowable.getMessage());
              }
          } else {
              response.setError(false);
          }
  
          return response;
      }
  
  
      private void waitResponse(Prompt prompt
          , ChatOptions options
          , AbstractBaseMessageResponse<?>[] messageResponse
          , CountDownLatch latch
          , Throwable[] failureThrowable) {
          chatStream(prompt, new StreamResponseListener() {
              @Override
              public void onMessage(ChatContext context, AiMessageResponse response) {
                  AiMessage message = response.getMessage();
                  if (message != null) message.setContent(message.getFullContent());
  
                  messageResponse[0] = response;
              }
  
              @Override
              public void onStop(ChatContext context) {
                  StreamResponseListener.super.onStop(context);
                  latch.countDown();
              }
  
              @Override
              public void onFailure(ChatContext context, Throwable throwable) {
                  logger.error(throwable.toString(), throwable);
                  failureThrowable[0] = throwable;
              }
          }, options);
  
          try {
              latch.await();
          } catch (InterruptedException e) {
              throw new RuntimeException(e);
          }
      }
  
  
      @Override
      public void chatStream(Prompt prompt, StreamResponseListener listener, ChatOptions options) {
          LlmClient llmClient = new WebSocketClient();
          String url = SparkLlmUtil.createURL(config);
          String payload = SparkLlmUtil.promptToPayload(prompt, config, options);
          LlmClientListener clientListener = new BaseLlmClientListener(this, llmClient, listener, prompt, aiMessageParser);
          llmClient.start(url, null, payload, clientListener, config);
      }
  
  
  }
  