Newer
Older
cortex-hub / ui / client-app / src / hooks / useCodeAssistant.js
import { useState, useEffect, useRef, useCallback } from "react";
import { getSessionId } from "../services/websocket";
import { getSessionTokenStatus, getSessionMessages, chatWithAI, getUserConfig, getSession } from "../services/apiService";

const useCodeAssistant = ({ pageContainerRef }) => {
  const [chatHistory, setChatHistory] = useState([]);
  const [isProcessing, setIsProcessing] = useState(false);
  const [errorMessage, setErrorMessage] = useState("");
  const [showErrorModal, setShowErrorModal] = useState(false);
  const [isConfigured, setIsConfigured] = useState(true);
  const [missingConfigs, setMissingConfigs] = useState([]);
  const [sessionId, setSessionId] = useState(null);
  const [tokenUsage, setTokenUsage] = useState({ token_count: 0, token_limit: 0, percentage: 0 });
  const [userConfigData, setUserConfigData] = useState(null);
  const [localActiveLLM, setLocalActiveLLM] = useState('');

  const sessionIdRef = useRef(null);
  const initialized = useRef(false);

  const fetchTokenUsage = useCallback(async () => {
    if (!sessionIdRef.current) return;
    try {
      const usage = await getSessionTokenStatus(sessionIdRef.current);
      setTokenUsage(usage);
    } catch (err) {
      console.warn("Failed to fetch token usage", err);
    }
  }, []);

  const fetchSessionHistory = useCallback(async (sid) => {
    try {
      const messagesData = await getSessionMessages(sid);
      if (messagesData && messagesData.messages) {
        const formattedHistory = messagesData.messages.map((msg) => ({
          isUser: msg.sender === "user",
          isPureAnswer: true,
          text: msg.content,
        }));
        setChatHistory(formattedHistory);
      }
    } catch (err) {
      console.warn("Failed to load chat history", err);
    }
  }, []);

  // Setup
  useEffect(() => {
    if (initialized.current) return;
    initialized.current = true;

    const setup = async () => {
      try {
        let configData = null;
        let provider = "gemini";
        try {
          configData = await getUserConfig();
          setUserConfigData(configData);
          if (configData.effective?.llm?.active_provider) {
            provider = configData.effective.llm.active_provider;
          }
        } catch (e) {
          console.warn("Could not load user config", e);
        }

        const sid = await getSessionId("coding_assistant", provider);
        setSessionId(sid);
        sessionIdRef.current = sid;

        let llm = provider;
        try {
          const sessionInfo = await getSession(sid);
          if (sessionInfo && sessionInfo.provider_name) {
            llm = sessionInfo.provider_name;
          }
        } catch (e) { console.warn("Could not check session provider", e); }
        setLocalActiveLLM(llm);

        // Config check
        const eff = configData?.effective || {};
        const missing = [];
        const llmProviders = eff.llm?.providers || {};
        const hasLLMKey = Object.values(llmProviders).some(p => p.api_key && p.api_key !== 'None');
        if (!hasLLMKey) missing.push("Language Model (LLM) API Key");

        if (missing.length > 0) {
          setIsConfigured(false);
          setMissingConfigs(missing);
        } else {
          setIsConfigured(true);
          setMissingConfigs([]);
        }

        await fetchSessionHistory(sid);
        await fetchTokenUsage();
      } catch (error) {
        console.error("Setup failed:", error);
      }
    };

    setup();
  }, [fetchSessionHistory, fetchTokenUsage]);

  const handleSendChat = useCallback(async (text) => {
    if (!isConfigured && text.trim().toLowerCase() !== "/new") {
      setErrorMessage("Coding Assistant requires a valid LLM configuration. Please visit Settings to set up your API keys.");
      setShowErrorModal(true);
      return;
    }

    if (text.trim().toLowerCase() === "/new") {
      setChatHistory([]);
      localStorage.removeItem("sessionId_coding_assistant");
      const newSid = await getSessionId("coding_assistant", localActiveLLM || "gemini");
      setSessionId(newSid);
      sessionIdRef.current = newSid;
      fetchTokenUsage();
      return;
    }

    setIsProcessing(true);
    setChatHistory((prev) => [...prev, { isUser: true, text }]);

    try {
      const response = await chatWithAI(sessionIdRef.current, text, localActiveLLM || "gemini");
      setChatHistory((prev) => [...prev, {
        isUser: false,
        isPureAnswer: true,
        text: response.answer,
        provider: response.provider_used
      }]);
      fetchTokenUsage();
    } catch (error) {
      setErrorMessage(error.message);
      setShowErrorModal(true);
    } finally {
      setIsProcessing(false);
    }
  }, [isConfigured, localActiveLLM, fetchTokenUsage]);

  const handleSwitchSession = useCallback(async (targetSessionId) => {
    localStorage.setItem("sessionId_coding_assistant", targetSessionId);
    setSessionId(targetSessionId);
    sessionIdRef.current = targetSessionId;
    setChatHistory([]);

    try {
      const sessionInfo = await getSession(targetSessionId);
      if (sessionInfo && sessionInfo.provider_name) {
        setLocalActiveLLM(sessionInfo.provider_name);
      }
      await fetchSessionHistory(targetSessionId);
      await fetchTokenUsage();
    } catch (error) {
      console.error("Failed to switch session:", error);
    }
  }, [fetchSessionHistory, fetchTokenUsage]);

  return {
    chatHistory,
    isProcessing,
    errorMessage,
    showErrorModal,
    tokenUsage,
    isConfigured,
    missingConfigs,
    handleSendChat,
    setShowErrorModal,
    handleSwitchSession,
    sessionId,
    userConfigData,
    localActiveLLM,
    setLocalActiveLLM
  };
};

export default useCodeAssistant;