File size: 2,198 Bytes
246d201
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import React from "react";
import { useDispatch, useSelector } from "react-redux";
import {
  useWsClient,
  WsClientProviderStatus,
} from "#/context/ws-client-provider";
import { createChatMessage } from "#/services/chat-service";
import { setCurrentAgentState } from "#/state/agent-slice";
import { addUserMessage } from "#/state/chat-slice";
import { clearFiles, clearInitialQuery } from "#/state/initial-query-slice";
import { RootState } from "#/store";
import { AgentState } from "#/types/agent-state";

export const useWSStatusChange = () => {
  const { send, status } = useWsClient();
  const { curAgentState } = useSelector((state: RootState) => state.agent);
  const dispatch = useDispatch();

  const statusRef = React.useRef<WsClientProviderStatus | null>(null);

  const { files, initialQuery } = useSelector(
    (state: RootState) => state.initialQuery,
  );

  const sendInitialQuery = (query: string, base64Files: string[]) => {
    const timestamp = new Date().toISOString();
    send(createChatMessage(query, base64Files, timestamp));
  };

  const dispatchInitialQuery = (query: string) => {
    sendInitialQuery(query, files);
    dispatch(clearFiles()); // reset selected files
    dispatch(clearInitialQuery()); // reset initial query
  };

  const handleAgentInit = () => {
    if (initialQuery) {
      dispatchInitialQuery(initialQuery);
    }
  };
  React.useEffect(() => {
    if (curAgentState === AgentState.INIT) {
      handleAgentInit();
    }
  }, [curAgentState]);

  React.useEffect(() => {
    if (statusRef.current === status) {
      return; // This is a check because of strict mode - if the status did not change, don't do anything
    }
    statusRef.current = status;

    if (status !== WsClientProviderStatus.DISCONNECTED && initialQuery) {
      dispatch(
        addUserMessage({
          content: initialQuery,
          imageUrls: files,
          timestamp: new Date().toISOString(),
          pending: true,
        }),
      );
    }

    if (status === WsClientProviderStatus.DISCONNECTED) {
      dispatch(setCurrentAgentState(AgentState.STOPPED));
    }
  }, [status]);
};