368 lines
10 KiB
TypeScript
368 lines
10 KiB
TypeScript
import { useCallback } from 'react'
|
|
|
|
import { Cortex } from '@cortexso/cortex.js'
|
|
import { Engine } from '@cortexso/cortex.js/resources'
|
|
import {
|
|
Assistant,
|
|
Model,
|
|
Message,
|
|
Thread,
|
|
ChatCompletionCreateParamsNonStreaming,
|
|
ChatCompletionCreateParamsStreaming,
|
|
AssistantCreateParams,
|
|
AssistantUpdateParams,
|
|
LlmEngine,
|
|
LlmEngines,
|
|
} from '@janhq/core'
|
|
|
|
import { useAtomValue } from 'jotai'
|
|
|
|
import { defaultThreadTitle } from '@/constants/Threads'
|
|
|
|
import { UpdateConfigMutationVariables } from './useEngineMutation'
|
|
import { MessageCreateMutationVariables } from './useMessageCreateMutation'
|
|
import { MessageDeleteMutationVariables } from './useMessageDeleteMutation'
|
|
import { MessageUpdateMutationVariables } from './useMessageUpdateMutation'
|
|
import { DownloadModelMutationVariables } from './useModelDownloadMutation'
|
|
|
|
import { hostAtom } from '@/helpers/atoms/AppConfig.atom'
|
|
|
|
const useCortex = () => {
|
|
const host = useAtomValue(hostAtom)
|
|
|
|
const cortex = new Cortex({
|
|
baseURL: host,
|
|
apiKey: '',
|
|
dangerouslyAllowBrowser: true,
|
|
})
|
|
|
|
const getEngineStatuses = useCallback(async (): Promise<Engine[]> => {
|
|
const engineResponse = await cortex.engines.list()
|
|
// @ts-expect-error incompatible types
|
|
const engineStatuses: Engine[] = engineResponse.body.data.map(
|
|
(engine: Engine) => {
|
|
return {
|
|
name: engine.name,
|
|
description: engine.description,
|
|
version: engine.version,
|
|
productName: engine.productName,
|
|
status: engine.status,
|
|
}
|
|
}
|
|
)
|
|
|
|
return engineStatuses
|
|
}, [cortex.engines])
|
|
|
|
const initializeEngine = useCallback(
|
|
async (engine: LlmEngine) => {
|
|
try {
|
|
await cortex.engines.init(engine)
|
|
} catch (err) {
|
|
console.error(err)
|
|
}
|
|
},
|
|
[cortex.engines]
|
|
)
|
|
|
|
const fetchAssistants = useCallback(async () => {
|
|
const assistants: Assistant[] = []
|
|
const response = await cortex.beta.assistants.list()
|
|
response.data.forEach((assistant) => {
|
|
assistants.push(assistant)
|
|
})
|
|
return assistants
|
|
}, [cortex.beta.assistants])
|
|
|
|
const fetchThreads = useCallback(async () => {
|
|
const threads: Thread[] = []
|
|
for await (const thread of cortex.beta.threads.list()) {
|
|
// @ts-expect-error each thread must have associated assistants
|
|
const assistants = thread['assistants'] as Assistant[]
|
|
if (!assistants || assistants.length === 0) continue
|
|
|
|
// @ts-expect-error each thread must have a title, else default to 'New Thread'
|
|
const title: string = thread['title'] ?? defaultThreadTitle
|
|
|
|
threads.push({
|
|
...thread,
|
|
title: title,
|
|
assistants: assistants,
|
|
})
|
|
}
|
|
return threads
|
|
}, [cortex.beta.threads])
|
|
|
|
const fetchModels = useCallback(async () => {
|
|
const models: Model[] = []
|
|
for await (const model of cortex.models.list()) {
|
|
// @ts-expect-error model should not be empty
|
|
const modelId = model.model
|
|
if (!modelId || modelId.length === 0) {
|
|
console.debug('Model id is empty, skipping', model)
|
|
continue
|
|
}
|
|
const engine = LlmEngines.find((engine) => engine === model.engine)
|
|
if (!engine) {
|
|
console.error(`Model ${modelId} has an invalid engine ${model.engine}`)
|
|
continue
|
|
}
|
|
|
|
models.push({
|
|
...model,
|
|
engine: engine,
|
|
model: modelId,
|
|
// @ts-expect-error each model must have associated files
|
|
files: model['files'],
|
|
})
|
|
}
|
|
return models
|
|
}, [cortex.models])
|
|
|
|
const fetchMessages = useCallback(
|
|
async (threadId: string) => {
|
|
try {
|
|
const messages: Message[] = []
|
|
const response = await cortex.beta.threads.messages.list(threadId)
|
|
response.data.forEach((message) => {
|
|
messages.push(message)
|
|
})
|
|
return messages
|
|
} catch (error) {
|
|
return []
|
|
}
|
|
},
|
|
[cortex.beta.threads.messages]
|
|
)
|
|
|
|
const startModel = useCallback(
|
|
async (modelId: string, options?: Record<string, unknown>) => {
|
|
await cortex.models.start(modelId, options ?? {})
|
|
},
|
|
[cortex.models]
|
|
)
|
|
|
|
const stopModel = useCallback(
|
|
async (modelId: string, options?: Record<string, unknown>) => {
|
|
await cortex.models.stop(modelId, options ?? {})
|
|
},
|
|
[cortex.models]
|
|
)
|
|
|
|
const chatCompletionNonStreaming = useCallback(
|
|
async (
|
|
chatCompletionCreateParams: ChatCompletionCreateParamsNonStreaming,
|
|
options?: Record<string, unknown>
|
|
// @ts-expect-error incompatible types
|
|
) => cortex.chat.completions.create(chatCompletionCreateParams, options),
|
|
[cortex.chat.completions]
|
|
)
|
|
|
|
const chatCompletionStreaming = useCallback(
|
|
async (
|
|
chatCompletionCreateParams: ChatCompletionCreateParamsStreaming,
|
|
options?: Record<string, unknown>
|
|
// @ts-expect-error incompatible types
|
|
) => cortex.chat.completions.create(chatCompletionCreateParams, options),
|
|
[cortex.chat.completions]
|
|
)
|
|
|
|
const deleteModel = useCallback(
|
|
async (modelId: string) => {
|
|
await cortex.models.del(modelId)
|
|
},
|
|
[cortex.models]
|
|
)
|
|
|
|
const cleanThread = useCallback(
|
|
async (threadId: string) => cortex.beta.threads.clean(threadId),
|
|
[cortex.beta.threads]
|
|
)
|
|
|
|
const deleteThread = useCallback(
|
|
async (threadId: string) => {
|
|
await cortex.beta.threads.del(threadId)
|
|
},
|
|
[cortex.beta.threads]
|
|
)
|
|
|
|
const updateThread = useCallback(
|
|
async (thread: Thread) => {
|
|
const result = await cortex.beta.threads.update(thread.id, thread)
|
|
console.debug(
|
|
`Update thread ${thread.id}, result: ${JSON.stringify(result, null, 2)}`
|
|
)
|
|
},
|
|
[cortex.beta.threads]
|
|
)
|
|
|
|
const deleteMessage = useCallback(
|
|
async (params: MessageDeleteMutationVariables) => {
|
|
const { threadId, messageId } = params
|
|
return cortex.beta.threads.messages.del(threadId, messageId)
|
|
},
|
|
[cortex.beta.threads]
|
|
)
|
|
|
|
const createMessage = useCallback(
|
|
async (params: MessageCreateMutationVariables) => {
|
|
const { threadId, createMessageParams } = params
|
|
return cortex.beta.threads.messages.create(threadId, createMessageParams)
|
|
},
|
|
[cortex.beta.threads]
|
|
)
|
|
|
|
const updateMessage = useCallback(
|
|
async (params: MessageUpdateMutationVariables) => {
|
|
const { threadId, messageId, data } = params
|
|
return cortex.beta.threads.messages.update(threadId, messageId, data)
|
|
},
|
|
[cortex.beta.threads]
|
|
)
|
|
|
|
const createThread = useCallback(
|
|
async (assistant: Assistant) => {
|
|
const createThreadResponse = await cortex.beta.threads.create({
|
|
// @ts-expect-error customize so that each thread will have an assistant
|
|
assistants: [assistant],
|
|
})
|
|
const thread: Thread = {
|
|
...createThreadResponse,
|
|
// @ts-expect-error each thread will have a title, else default to 'New Thread'
|
|
title: createThreadResponse.title ?? defaultThreadTitle,
|
|
assistants: [assistant],
|
|
}
|
|
return thread
|
|
},
|
|
[cortex.beta.threads]
|
|
)
|
|
|
|
const updateModel = useCallback(
|
|
async (modelId: string, options: Record<string, unknown>) => {
|
|
try {
|
|
return await cortex.models.update(modelId, options)
|
|
} catch (err) {
|
|
console.error(err)
|
|
}
|
|
},
|
|
[cortex.models]
|
|
)
|
|
|
|
const downloadModel = useCallback(
|
|
async (variables: DownloadModelMutationVariables) => {
|
|
const { modelId, fileName, persistedModelId } = variables
|
|
const response = await fetch(`${host}/models/${modelId}/pull`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'accept': 'application/json',
|
|
// eslint-disable-next-line @typescript-eslint/naming-convention
|
|
'Content-Type': 'application/json',
|
|
},
|
|
body: JSON.stringify({
|
|
fileName: fileName,
|
|
persistedModelId: persistedModelId,
|
|
}),
|
|
})
|
|
if (!response.ok) {
|
|
const responseJson = await response.json()
|
|
const errorMessage: string =
|
|
(responseJson.error?.message as string) ??
|
|
`Failed to download model ${modelId}`
|
|
throw new Error(errorMessage)
|
|
}
|
|
},
|
|
[host]
|
|
)
|
|
|
|
const abortDownload = useCallback(
|
|
async (downloadId: string) => {
|
|
try {
|
|
return await cortex.models.abortDownload(downloadId)
|
|
} catch (err) {
|
|
console.error(err)
|
|
}
|
|
},
|
|
[cortex.models]
|
|
)
|
|
|
|
const createAssistant = useCallback(
|
|
async (createParams: AssistantCreateParams) =>
|
|
cortex.beta.assistants.create(createParams),
|
|
[cortex.beta.assistants]
|
|
)
|
|
|
|
const updateAssistant = useCallback(
|
|
async (assistantId: string, updateParams: AssistantUpdateParams) =>
|
|
cortex.beta.assistants.update(assistantId, updateParams),
|
|
[cortex.beta.assistants]
|
|
)
|
|
|
|
// TODO: add this to cortex-node
|
|
const registerEngineConfig = useCallback(
|
|
async (variables: UpdateConfigMutationVariables) => {
|
|
const { engine, config } = variables
|
|
try {
|
|
await cortex.engines.update(engine, config)
|
|
} catch (err) {
|
|
console.error(err)
|
|
}
|
|
},
|
|
[cortex.engines]
|
|
)
|
|
|
|
// add this to cortex-node?
|
|
const createModel = useCallback(
|
|
(model: Model) =>
|
|
fetch(`${host}/models`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'accept': 'application/json',
|
|
// eslint-disable-next-line @typescript-eslint/naming-convention
|
|
'Content-Type': 'application/json',
|
|
},
|
|
body: JSON.stringify(model),
|
|
}),
|
|
[host]
|
|
)
|
|
|
|
const isSystemAlive = useCallback(async () => {
|
|
try {
|
|
await cortex.system.status()
|
|
return true
|
|
} catch {
|
|
return false
|
|
}
|
|
}, [cortex.system])
|
|
|
|
return {
|
|
fetchAssistants,
|
|
fetchThreads,
|
|
fetchModels,
|
|
fetchMessages,
|
|
startModel,
|
|
stopModel,
|
|
chatCompletionStreaming,
|
|
deleteModel,
|
|
deleteThread,
|
|
deleteMessage,
|
|
cleanThread,
|
|
updateThread,
|
|
createMessage,
|
|
updateMessage,
|
|
createThread,
|
|
downloadModel,
|
|
abortDownload,
|
|
createAssistant,
|
|
updateAssistant,
|
|
updateModel,
|
|
chatCompletionNonStreaming,
|
|
registerEngineConfig,
|
|
createModel,
|
|
initializeEngine,
|
|
getEngineStatuses,
|
|
isSystemAlive,
|
|
}
|
|
}
|
|
|
|
export default useCortex
|