74 lines
1.9 KiB
TypeScript
74 lines
1.9 KiB
TypeScript
|
import { detectThinkKeyword, extractPromptText } from "./detector"
|
||
|
|
import { getHighVariant, isAlreadyHighVariant } from "./switcher"
|
||
|
|
import type { ThinkModeState, ThinkModeInput } from "./types"
|
||
|
|
|
||
|
|
export * from "./detector"
|
||
|
|
export * from "./switcher"
|
||
|
|
export * from "./types"
|
||
|
|
|
||
|
|
const thinkModeState = new Map<string, ThinkModeState>()
|
||
|
|
|
||
|
|
export function clearThinkModeState(sessionID: string): void {
|
||
|
|
thinkModeState.delete(sessionID)
|
||
|
|
}
|
||
|
|
|
||
|
|
export function createThinkModeHook() {
|
||
|
|
return {
|
||
|
|
"chat.params": async (
|
||
|
|
output: ThinkModeInput,
|
||
|
|
sessionID: string
|
||
|
|
): Promise<void> => {
|
||
|
|
const promptText = extractPromptText(output.parts)
|
||
|
|
|
||
|
|
const state: ThinkModeState = {
|
||
|
|
requested: false,
|
||
|
|
modelSwitched: false,
|
||
|
|
}
|
||
|
|
|
||
|
|
if (!detectThinkKeyword(promptText)) {
|
||
|
|
thinkModeState.set(sessionID, state)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
state.requested = true
|
||
|
|
|
||
|
|
const currentModel = output.message.model
|
||
|
|
if (!currentModel) {
|
||
|
|
thinkModeState.set(sessionID, state)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
state.providerID = currentModel.providerID
|
||
|
|
state.modelID = currentModel.modelID
|
||
|
|
|
||
|
|
if (isAlreadyHighVariant(currentModel.modelID)) {
|
||
|
|
thinkModeState.set(sessionID, state)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
const highVariant = getHighVariant(currentModel.modelID)
|
||
|
|
|
||
|
|
if (!highVariant) {
|
||
|
|
thinkModeState.set(sessionID, state)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
output.message.model = {
|
||
|
|
providerID: currentModel.providerID,
|
||
|
|
modelID: highVariant,
|
||
|
|
}
|
||
|
|
state.modelSwitched = true
|
||
|
|
thinkModeState.set(sessionID, state)
|
||
|
|
},
|
||
|
|
|
||
|
|
event: async ({ event }: { event: { type: string; properties?: unknown } }) => {
|
||
|
|
if (event.type === "session.deleted") {
|
||
|
|
const props = event.properties as { info?: { id?: string } } | undefined
|
||
|
|
if (props?.info?.id) {
|
||
|
|
thinkModeState.delete(props.info.id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
}
|
||
|
|
}
|