46 lines
1.7 KiB
TypeScript
46 lines
1.7 KiB
TypeScript
import type { ComarkPlugin } from "comark";
|
|
import type { SearchSource, StoredToolValue } from "./tool-rendering";
|
|
import { parseToolValue, readSearchSources } from "./tool-rendering";
|
|
|
|
export const citationPlugin: ComarkPlugin = {
|
|
name: "source-citations",
|
|
markdownItPlugins: [
|
|
(markdown) => {
|
|
markdown.inline.ruler.before("link", "source_citation", (state, silent) => {
|
|
if (state.linkLevel > 0 || !state.src.startsWith("[[cite:", state.pos)) {
|
|
return false;
|
|
}
|
|
const match = /^\[\[cite:([a-zA-Z0-9_-]+(?:,[a-zA-Z0-9_-]+)*)\]\]/.exec(state.src.slice(state.pos));
|
|
if (!match) {
|
|
return false;
|
|
}
|
|
if (silent) {
|
|
return false;
|
|
}
|
|
const token = state.push("mdc_inline_component", "source-cite", 0);
|
|
token.attrSet("ids", match[1] ?? "");
|
|
state.pos += match[0].length;
|
|
return true;
|
|
});
|
|
},
|
|
],
|
|
};
|
|
|
|
export const collectCitationSources = (
|
|
parts: readonly { toolCall?: { toolName: string; output?: StoredToolValue } | null }[],
|
|
): ReadonlyMap<string, SearchSource> => {
|
|
const sources = new Map<string, SearchSource>();
|
|
for (const part of parts) {
|
|
if (part.toolCall?.toolName !== "search") {
|
|
continue;
|
|
}
|
|
const results = readSearchSources(parseToolValue(part.toolCall.output), { deduplicate: false });
|
|
for (const source of results ?? []) {
|
|
if (source.sourceId) {
|
|
sources.set(source.sourceId, source);
|
|
}
|
|
}
|
|
}
|
|
return sources;
|
|
};
|