Skip to content
Source
Sign inSign up

composerModelMentions.ts

104 lines4.0 KB

1import { rankModelPickerSearch, type ModelPickerSearchableModel } from "./modelSearch.ts";
2​
3const COMPOSER_MODEL_MENTION_LIMIT = 8;
4​
5const escapeRegExp = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/gu, "\\$&");
6​
7/** Mentions can't contain spaces, so `:`, `/` and `-` split words, as in `@gpt:traex`. */
8function mentionSearchQueries(query: string): ReadonlyArray<string> {
9 const spaced = query.replace(/[:/]+/gu, " ").trim();
10 const unhyphenated = spaced.replace(/-+/gu, " ").trim();
11 // One letter would fill the menu with models and push files out of it.
12 if (unhyphenated.replace(/\s+/gu, "").length < 2) return [];
13 return unhyphenated === spaced ? [spaced] : [spaced, unhyphenated];
15​
16/** Returns the models matching an `@` mention, best first. A bare `@` matches none. */
17export function rankComposerModelMentions<T>(
18 candidates: ReadonlyArray<T>,
19 query: string,
20 toSearchable: (candidate: T) => ModelPickerSearchableModel,
21): T[] {
22 for (const words of mentionSearchQueries(query)) {
23 // Every word must start a word of the model's names. Fuzzy hits would push files out of the `@` menu.
24 const tokens = words
25 .toLowerCase()
26 .split(/\s+/u)
27 .map((token) => new RegExp(`(?:^|[^\\p{L}\\p{N}])${escapeRegExp(token)}`, "u"));
28 const literal = candidates.filter((candidate) => {
29 const model = toSearchable(candidate);
30 const haystack = [
31 model.name,
32 model.shortName,
33 model.subProvider,
34 model.driverKind,
35 model.providerDisplayName,
36 ]
37 .join(" ")
38 .toLowerCase();
39 return tokens.every((token) => token.test(haystack));
40 });
41 const ranked = rankModelPickerSearch(literal, words, toSearchable);
42 if (ranked.length > 0) return ranked.slice(0, COMPOSER_MODEL_MENTION_LIMIT);
43 }
44 return [];
46​
47/** Prompt text naming a provider instance and model for the agent. No `@`, which the editor would turn into a file chip. */
48export function formatComposerModelMention(instanceId: string, model: string): string {
49 return `model:${instanceId}/${model}`;
51​
52export interface ComposerModelMentionOccurrence {
53 readonly instanceId: string;
54 readonly model: string;
55 readonly source: string;
56 readonly start: number;
57 readonly end: number;
59​
60// Instance ids match ProviderInstanceId; model slugs may contain `/`, as in `opencode/anthropic/opus`.
61const MODEL_MENTION_PATTERN = /(?<![^\s,(["'])model:([a-zA-Z][a-zA-Z0-9_-]*)\/([^\s,"']+)/gu;
62const SENTENCE_PUNCTUATION = ".;:!?";
63const CLOSING_BRACKETS: Readonly<Record<string, string>> = { ")": "(", "]": "[" };
64​
65function countOf(text: string, char: string): number {
66 return text.split(char).length - 1;
68​
69// Slugs like `foo(bar)` or `GPT-6.1-Sol[1m]` keep balanced brackets; prose wrapping them does not.
70function trimModelSlug(slug: string): string {
71 let end = slug.length;
72 const unmatched = new Map<string, number>();
73 while (end > 0) {
74 const last = slug[end - 1]!;
75 if (SENTENCE_PUNCTUATION.includes(last)) {
76 end -= 1;
77 continue;
78 }
79 const opening = CLOSING_BRACKETS[last];
80 if (opening === undefined) break;
81 const excess = unmatched.get(last) ?? countOf(slug, last) - countOf(slug, opening);
82 if (excess <= 0) break;
83 unmatched.set(last, excess - 1);
84 end -= 1;
85 }
86 return slug.slice(0, end);
88​
89/** Finds the `model:<instance>/<model>` references that `formatComposerModelMention` writes. */
90export function collectComposerModelMentions(
91 text: string,
92): ReadonlyArray<ComposerModelMentionOccurrence> {
93 if (!text.includes("model:")) return [];
94 const occurrences: ComposerModelMentionOccurrence[] = [];
95 for (const match of text.matchAll(MODEL_MENTION_PATTERN)) {
96 const instanceId = match[1]!;
97 const model = trimModelSlug(match[2]!);
98 if (model.length === 0) continue;
99 const start = match.index;
100 const end = start + "model:".length + instanceId.length + 1 + model.length;
101 occurrences.push({ instanceId, model, source: text.slice(start, end), start, end });
102 }
103 return occurrences;