LocalMind/hooks/useLMStudio.ts

import { useCallback, useRef } from 'react';
import { SYSTEM_PROMPT } from '../constants/SystemPrompt';
import * as api from '../services/lmStudioApi';
import { ChatMessage as ApiMessage } from '../services/types';
import { ChatMessage, useAppStore } from '../store/useAppStore';

const NO_MODEL_MESSAGE = '⚠️ No model selected. Please select a model first.';

function buildApiMessages(messages: ChatMessage[]): ApiMessage[] {
    return [
        { role: 'system', content: SYSTEM_PROMPT },
        ...messages.map((m) => ({ role: m.role, content: m.content })),
    ];
}

export function useLMStudio() {
    const serverConfig = useAppStore((s) => s.serverConfig);
    const isConnected = useAppStore((s) => s.isConnected);
    const setIsConnected = useAppStore((s) => s.setIsConnected);
    const availableModels = useAppStore((s) => s.availableModels);
    const setAvailableModels = useAppStore((s) => s.setAvailableModels);
    const activeModelId = useAppStore((s) => s.activeModelId);
    const setActiveModelId = useAppStore((s) => s.setActiveModelId);
    const messages = useAppStore((s) => s.messages);
    const addMessage = useAppStore((s) => s.addMessage);
    const appendToMessage = useAppStore((s) => s.appendToMessage);
    const isStreaming = useAppStore((s) => s.isStreaming);
    const setIsStreaming = useAppStore((s) => s.setIsStreaming);

    const abortRef = useRef<{ abort: () => void } | null>(null);

    const checkConnection = useCallback(async () => {
        const ok = await api.testConnection(serverConfig);
        setIsConnected(ok);
        return ok;
    }, [serverConfig, setIsConnected]);

    const fetchModels = useCallback(async () => {
        try {
            const currentConfig = useAppStore.getState().serverConfig;
            const models = await api.getModels(currentConfig);
            setAvailableModels(models);
            setIsConnected(true);

            if (!activeModelId && models.length > 0) {
                setActiveModelId(models[0].id);
            }
            return models;
        } catch {
            setIsConnected(false);
            setAvailableModels([]);
            return [];
        }
    }, [activeModelId, setAvailableModels, setIsConnected, setActiveModelId]);

    // Adds an empty assistant message and streams the model's answer into it.
    const streamAnswer = useCallback((modelId: string, apiMessages: ApiMessage[]) => {
        const assistantId = addMessage({ role: 'assistant', content: '' });
        setIsStreaming(true);

        const finish = () => {
            setIsStreaming(false);
            abortRef.current = null;
            useAppStore.getState().saveCurrentChat();
        };

        try {
            abortRef.current = api.chatCompletionStream(
                serverConfig,
                modelId,
                apiMessages,
                (chunk) => appendToMessage(assistantId, chunk),
                finish,
                (error) => {
                    appendToMessage(assistantId, `

⚠️ Error: ${error.message}`);
                    finish();
                },
            );
        } catch (err) {
            appendToMessage(assistantId, `⚠️ Connection failed: ${(err as Error).message}`);
            finish();
        }
    }, [serverConfig, addMessage, appendToMessage, setIsStreaming]);

    const sendMessage = useCallback(async (text: string) => {
        if (isStreaming) return;

        addMessage({ role: 'user', content: text });

        const modelId = useAppStore.getState().activeModelId;
        if (!modelId) {
            addMessage({ role: 'assistant', content: NO_MODEL_MESSAGE });
            return;
        }

        streamAnswer(modelId, buildApiMessages(useAppStore.getState().messages));
    }, [isStreaming, addMessage, streamAnswer]);

    const regenerateResponse = useCallback(async () => {
        if (isStreaming) return;

        const { messages: msgs, activeModelId: modelId, deleteMessagesAfter } = useAppStore.getState();
        const lastUserMsg = [...msgs].reverse().find((m) => m.role === 'user');
        if (!lastUserMsg) return;

        // Checked before deleting anything, so the current answer is not lost for nothing.
        if (!modelId) {
            addMessage({ role: 'assistant', content: NO_MODEL_MESSAGE });
            return;
        }

        // Drop the previous answer(s) and ask again from the last user message.
        deleteMessagesAfter(lastUserMsg.id);
        streamAnswer(modelId, buildApiMessages(useAppStore.getState().messages));
    }, [isStreaming, addMessage, streamAnswer]);

    const stopStreaming = useCallback(() => {
        if (abortRef.current) {
            abortRef.current.abort();
            abortRef.current = null;
            setIsStreaming(false);
            useAppStore.getState().saveCurrentChat();
        }
    }, [setIsStreaming]);

    return {
        isConnected,
        isStreaming,
        availableModels,
        activeModelId,
        messages,
        sendMessage,
        stopStreaming,
        fetchModels,
        checkConnection,
        regenerateResponse,
    };
}