62 lines
1.4 KiB
TypeScript
62 lines
1.4 KiB
TypeScript
|
import type { ContextCollector } from "./collector"
|
||
|
|
|
||
|
|
const MESSAGE_SEPARATOR = "\n\n---\n\n"
|
||
|
|
|
||
|
|
interface OutputPart {
|
||
|
|
type: string
|
||
|
|
text?: string
|
||
|
|
[key: string]: unknown
|
||
|
|
}
|
||
|
|
|
||
|
|
interface InjectionResult {
|
||
|
|
injected: boolean
|
||
|
|
contextLength: number
|
||
|
|
}
|
||
|
|
|
||
|
|
export function injectPendingContext(
|
||
|
|
collector: ContextCollector,
|
||
|
|
sessionID: string,
|
||
|
|
parts: OutputPart[]
|
||
|
|
): InjectionResult {
|
||
|
|
if (!collector.hasPending(sessionID)) {
|
||
|
|
return { injected: false, contextLength: 0 }
|
||
|
|
}
|
||
|
|
|
||
|
|
const textPartIndex = parts.findIndex((p) => p.type === "text" && p.text !== undefined)
|
||
|
|
if (textPartIndex === -1) {
|
||
|
|
return { injected: false, contextLength: 0 }
|
||
|
|
}
|
||
|
|
|
||
|
|
const pending = collector.consume(sessionID)
|
||
|
|
const originalText = parts[textPartIndex].text ?? ""
|
||
|
|
parts[textPartIndex].text = `${pending.merged}${MESSAGE_SEPARATOR}${originalText}`
|
||
|
|
|
||
|
|
return {
|
||
|
|
injected: true,
|
||
|
|
contextLength: pending.merged.length,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
interface ChatMessageInput {
|
||
|
|
sessionID: string
|
||
|
|
agent?: string
|
||
|
|
model?: { providerID: string; modelID: string }
|
||
|
|
messageID?: string
|
||
|
|
}
|
||
|
|
|
||
|
|
interface ChatMessageOutput {
|
||
|
|
message: Record<string, unknown>
|
||
|
|
parts: OutputPart[]
|
||
|
|
}
|
||
|
|
|
||
|
|
export function createContextInjectorHook(collector: ContextCollector) {
|
||
|
|
return {
|
||
|
|
"chat.message": async (
|
||
|
|
input: ChatMessageInput,
|
||
|
|
output: ChatMessageOutput
|
||
|
|
): Promise<void> => {
|
||
|
|
injectPendingContext(collector, input.sessionID, output.parts)
|
||
|
|
},
|
||
|
|
}
|
||
|
|
}
|