1import { collectComposerModelMentions } from "@t3tools/client-runtime/composer-model-mentions";
2import type { AssistantCitation } from "@t3tools/contracts";
3import { collectAssistantCitations } from "@t3tools/shared/assistantCitations";
4import { collectComposerContextReferences } from "@t3tools/shared/composerContextReferences";
5import {
6 collectComposerInlineTokens,
7 composerMathRanges,
8 type ComposerInlineToken,
9} from "@t3tools/shared/composerInlineTokens";
10
11export type ComposerPromptSegment =
12 | {
13 type: "text";
14 text: string;
15 }
16 | {
17 type: "mention";
18 path: string;
19 source: string;
20 }
21 | {
22 type: "skill";
23 name: string;
24 source: string;
25 }
26 | {
27 type: "model";
28 instanceId: string;
29 model: string;
30 source: string;
31 }
32 | {
33 type: "citation";
34 citation: AssistantCitation;
35 source: string;
36 }
37 | {
38 type: "context-reference";
39 kind: string;
40 contextId: string;
41 label: string;
42 source: string;
43 };
44
45function rangeIncludesIndex(start: number, end: number, index: number): boolean {
46 return start <= index && index < end;
47}
48
49function pushTextSegment(segments: ComposerPromptSegment[], text: string): void {
50 if (!text) return;
51 const last = segments[segments.length - 1];
52 if (last && last.type === "text") {
53 last.text += text;
54 return;
55 }
56 segments.push({ type: "text", text });
57}
58
59function forEachPromptTextSlice(
60 prompt: string,
61 visitor: (text: string, promptOffset: number) => boolean | void,
62): boolean {
63 return prompt.length > 0 && visitor(prompt, 0) === true;
64}
65
66function forEachMentionMatch(
67 prompt: string,
68 visitor: (
69 match: Extract<ComposerInlineToken, { type: "mention" }>,
70 promptOffset: number,
71 ) => boolean | void,
72): boolean {
73 return forEachPromptTextSlice(prompt, (text, promptOffset) => {
74 for (const match of collectComposerPromptInlineTokens(text)) {
75 if (match.type !== "mention") {
76 continue;
77 }
78 if (visitor(match, promptOffset) === true) {
79 return true;
80 }
81 }
82 return false;
83 });
84}
85
86// Like `$skill`, a typed reference waits for whitespace so `gpt-6` cannot chip while typing `gpt-6.1`.
87function collectModelTokens(text: string, inlineTokens: ReadonlyArray<ComposerInlineToken>) {
88 const mentions = collectComposerModelMentions(text).filter((mention) =>
89 /\s/.test(text[mention.end] ?? ""),
90 );
91 if (mentions.length === 0) return [];
92 const taken = [...inlineTokens, ...composerMathRanges(text)];
93 return mentions
94 .filter(
95 (mention) => !taken.some((range) => mention.start < range.end && mention.end > range.start),
96 )
97 .map((mention) => ({ ...mention, type: "model" as const }));
98}
99
100export function collectComposerPromptInlineTokens(text: string) {
101 const inlineTokens = collectComposerInlineTokens(text);
102 const models = collectModelTokens(text, inlineTokens);
103 const tokens =
104 models.length === 0
105 ? inlineTokens
106 : [...inlineTokens, ...models].sort((left, right) => left.start - right.start);
107 const citations = collectAssistantCitations(text);
108 const references = collectComposerContextReferences(text);
109 if (citations.length === 0 && references.length === 0) return tokens;
110
111 // An unfinished @ mention can otherwise consume the start of a link label.
112 const links = [
113 ...citations.map((match) => ({ ...match, type: "citation" as const })),
114 ...references.map((match) => ({ ...match, type: "context-reference" as const })),
115 ];
116 return [
117 ...tokens.filter(
118 (token) => !links.some((link) => token.start < link.end && token.end > link.start),
119 ),
120 ...links,
121 ].sort((left, right) => left.start - right.start);
122}
123
124function splitPromptTextIntoComposerSegments(text: string): ComposerPromptSegment[] {
125 const segments: ComposerPromptSegment[] = [];
126 if (!text) {
127 return segments;
128 }
129
130 const tokenMatches = collectComposerPromptInlineTokens(text);
131 let cursor = 0;
132 for (const match of tokenMatches) {
133 if (match.start < cursor) {
134 continue;
135 }
136
137 if (match.start > cursor) {
138 pushTextSegment(segments, text.slice(cursor, match.start));
139 }
140
141 if (match.type === "citation") {
142 segments.push({ type: "citation", citation: match.citation, source: match.source });
143 } else if (match.type === "context-reference") {
144 segments.push({
145 type: "context-reference",
146 kind: match.kind,
147 contextId: match.contextId,
148 label: match.label,
149 source: match.source,
150 });
151 } else if (match.type === "mention") {
152 segments.push({
153 type: "mention",
154 path: match.value,
155 source: match.source,
156 });
157 } else if (match.type === "model") {
158 segments.push({
159 type: "model",
160 instanceId: match.instanceId,
161 model: match.model,
162 source: match.source,
163 });
164 } else {
165 segments.push({ type: "skill", name: match.value, source: match.source });
166 }
167
168 cursor = match.end;
169 }
170
171 if (cursor < text.length) {
172 pushTextSegment(segments, text.slice(cursor));
173 }
174
175 return segments;
176}
177
178export function selectionTouchesMentionBoundary(
179 prompt: string,
180 start: number,
181 end: number,
182): boolean {
183 if (!prompt || start >= end) {
184 return false;
185 }
186
187 return forEachMentionMatch(prompt, (match, promptOffset) => {
188 const mentionStart = promptOffset + match.start;
189 const mentionEnd = promptOffset + match.end;
190 const beforeMentionIndex = mentionStart - 1;
191 const afterMentionIndex = mentionEnd;
192
193 if (
194 beforeMentionIndex >= 0 &&
195 /\s/.test(prompt[beforeMentionIndex] ?? "") &&
196 rangeIncludesIndex(start, end, beforeMentionIndex)
197 ) {
198 return true;
199 }
200
201 if (
202 afterMentionIndex < prompt.length &&
203 /\s/.test(prompt[afterMentionIndex] ?? "") &&
204 rangeIncludesIndex(start, end, afterMentionIndex)
205 ) {
206 return true;
207 }
208 return false;
209 });
210}
211
212export function splitPromptIntoComposerSegments(prompt: string): ComposerPromptSegment[] {
213 return splitPromptTextIntoComposerSegments(prompt);
214}