Industrial instruction transcription interaction protocol (WebSocket)

package nls.alibaba;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.codec.binary.Hex;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.nio.file.Files;
import java.security.SecureRandom;
import java.util.HashMap;
import java.util.Map;

import static java.lang.Thread.sleep;

public class TingwuIndustrialInstructionUsage extends WebSocketClient {

    private final byte[] audioData;

    private final String taskId;

    private final ObjectMapper objectMapper = new ObjectMapper();

    private boolean audioSent = false;

    private static final int CHUNK_SIZE = 3200;

    private static final long SEND_INTERVAL_MS = 100; // Simulate streaming send interval

    private static final String apiKey = "Replace with your API key";

    public TingwuIndustrialInstructionUsage(URI serverUri, String audioFilePath) throws IOException {
        super(serverUri, createHeaders());
        this.audioData = Files.readAllBytes(new File(audioFilePath).toPath());

        // Generate a random 16-character string
        SecureRandom random = new SecureRandom();
        byte[] taskIdBytes = new byte[8];
        random.nextBytes(taskIdBytes);
        this.taskId = Hex.encodeHexString(taskIdBytes);
    }

    private static Map<String, String> createHeaders() {
        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", "Bearer " + apiKey);
        headers.put("Upgrade", "websocket");
        headers.put("Connection", "Upgrade");
        return headers;
    }

    @Override
    public void onOpen(ServerHandshake serverHandshake) {
        System.out.println("WebSocket connected.");
        System.out.println("Task ID: " + taskId);

        // Construct the run-task message
        String runTaskMessage = "{\n" +
                "  \"header\": {\n" +
                "    \"action\": \"run-task\",\n" +
                "    \"task_id\": \"" + taskId + "\",\n" +
                "    \"streaming\": \"duplex\"\n" +
                "  },\n" +
                "  \"payload\": {\n" +
                "    \"model\": \"tingwu-industrial-instruction\",\n" +
                "    \"task_group\": \"aigc\",\n" +
                "    \"task\": \"multimodal-generation\",\n" +
                "    \"function\": \"generation\",\n" +
                "    \"input\": {\n" +
                "      \"appId\": \"Replace with your application ID\",\n" +
                "      \"directive\": \"start\"\n" +
                "    },\n" +
                "    \"parameters\": {\n" +
                "      \"sampleRate\": 16000,\n" +
                "      \"format\": \"pcm\",\n" +
                "      \"maxEndSilence\": 1500,\n" +
                "      \"terminology\": \"Replace with your instruction set ID\"\n" +
                "    }\n" +
                "  }\n" +
                "}";

        System.out.println("Sending run-task: " + runTaskMessage);
        send(runTaskMessage);
    }

    @Override
    public void onMessage(String message) {
        System.out.println("Received text message: " + message);

        try {
            JsonNode root = objectMapper.readTree(message);

            JsonNode header = root.get("header");
            JsonNode payload = root.get("payload");

            if (header != null) {
                String event = header.path("event").asText("");
                // This message can be ignored
                if ("task-started".equals(event)) {
                    return;
                }

                String action = payload.path("output").path("action").asText("");
                switch (action) {
                    case "speech-listen":
                        System.out.println("Received speech-listen event, sending full audio...");
                        sendAudioFile();
                        break;
                    case "task-failed":
                        String errorCode = payload.path("output").path("errorCode").asText("N/A");
                        String errorMsg = payload.path("output").path("errorMessage").asText("N/A");
                        System.out.println("Task failed: " + errorCode + " - " + errorMsg);
                        close();
                        break;
                    case "recognize-result":
                        JsonNode output = payload.get("output");
                        System.out.println("Recognition result: " + (output != null ? output.toString() : "N/A"));
                        break;
                    case "ai-result":
                        JsonNode correction = payload.path("output").path("aiResult").path("correction");
                        System.out.println("AI correction result: " + correction.asText("N/A"));
                        break;
                    case "speech-end":
                        System.out.println("Speech ended. Closing...");
                        close();
                        break;
                    default:
                        System.out.println("Unknown message received: " + message);
                        break;
                }
            } else {
                System.out.println("Unknown message format (no header): " + message);
            }

        } catch (JsonProcessingException e) {
            System.err.println("Invalid JSON received: " + e.getMessage());
            System.out.println("Raw message: " + message);
        }
    }

    @Override
    public void onClose(int code, String reason, boolean b) {
        System.out.println("WebSocket closed. Code: " + code + ", Reason: " + reason);
    }

    @Override
    public void onError(Exception ex) {
        System.err.println("WebSocket error: " + ex.getMessage());
        ex.printStackTrace();
    }

    private void sendAudioFile() {
        if (audioSent) {
            System.out.println("Audio already sent, skipping...");
            return;
        }
        audioSent = true;

        int offset = 0;
        while (offset < audioData.length) {
            int len = Math.min(CHUNK_SIZE, audioData.length - offset);
            byte[] chunk = new byte[len];
            System.arraycopy(audioData, offset, chunk, 0, len);
            try {
                send(chunk);
                offset += len;
                System.out.println("Sent chunk: " + offset + "/" + audioData.length);
                sleep(100);
            } catch (Exception e) {
                System.err.println("Send failed at offset " + offset);
                e.printStackTrace();
                close();
                return;
            }
        }
        System.out.println("Audio sent completely.");
        sendFinishTask();

    }

    // Send the finish-task message
    private void sendFinishTask() {
        String finishTaskMessage = "{\n" +
                "  \"header\": {\n" +
                "    \"action\": \"finish-task\",\n" +
                "    \"task_id\": \"" + taskId + "\",\n" +
                "    \"streaming\": \"duplex\"\n" +
                "  },\n" +
                "  \"payload\": {\n" +
                "    \"model\": \"tingwu-industrial-instruction\",\n" +
                "    \"task_group\": \"aigc\",\n" +
                "    \"task\": \"multimodal-generation\",\n" +
                "    \"function\": \"generation\",\n" +
                "    \"input\": {\n" +
                "      \"directive\": \"stop\"\n" +
                "    }\n" +
                "  }\n" +
                "}";

        System.out.println("Sending finish-task: " + finishTaskMessage);
        send(finishTaskMessage);
    }

    public static void main(String[] args) {
        try {
            URI wsUri = new URI("wss://dashscope.aliyuncs.com/api-ws/v1/inference");
            String audioPath = "Replace with your audio file path";
            TingwuIndustrialInstructionUsage client = new TingwuIndustrialInstructionUsage(wsUri, audioPath);
            client.connect();
        } catch (Exception e) {
            System.err.println("Initialization error: " + e.getMessage());
            e.printStackTrace();
        }
    }
}