Skip to content

Commit 20ab6ad

Browse files
nachawatinachawati
committed
feat(vscode): inject workspace context into chat, add latex math rendering, listClasses lsp endpoint
1 parent e5ae5e8 commit 20ab6ad

3 files changed

Lines changed: 195 additions & 30 deletions

File tree

‎packages/lsp/src/browserServerMain.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3390,5 +3390,32 @@ connection.onRequest(
33903390
},
33913391
);
33923392

3393+
// List all top-level classes across all loaded documents (for AI chat context)
3394+
connection.onRequest("modelscript/listClasses", (): { classes: { name: string; kind: string; uri: string }[] } => {
3395+
const classes: { name: string; kind: string; uri: string }[] = [];
3396+
const seen = new Set<string>();
3397+
3398+
for (const [uri, ctx] of documentContexts.entries()) {
3399+
try {
3400+
for (const element of ctx.elements) {
3401+
if (element instanceof ModelicaClassInstance && element.name) {
3402+
if (!seen.has(element.name)) {
3403+
seen.add(element.name);
3404+
classes.push({
3405+
name: element.name,
3406+
kind: element.classKind ?? "class",
3407+
uri,
3408+
});
3409+
}
3410+
}
3411+
}
3412+
} catch {
3413+
// Skip problematic contexts
3414+
}
3415+
}
3416+
3417+
return { classes };
3418+
});
3419+
33933420
// Listen on the connection
33943421
connection.listen();

‎packages/vscode/src/chatPanel.ts‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,25 @@ export class ChatPanel {
4949
case "getActiveFileContext":
5050
this.sendActiveFileContext();
5151
break;
52+
case "listClasses":
53+
await this.handleListClasses(msg);
54+
break;
5255
}
5356
},
5457
null,
5558
this.disposables,
5659
);
5760

61+
// Auto-send active file context when panel opens
62+
this.sendActiveFileContext();
63+
64+
// Re-send active file context when user switches editors
65+
this.disposables.push(
66+
vscode.window.onDidChangeActiveTextEditor(() => {
67+
this.sendActiveFileContext();
68+
}),
69+
);
70+
5871
this.panel.onDidDispose(() => this.dispose(), null, this.disposables);
5972
}
6073

@@ -87,8 +100,28 @@ export class ChatPanel {
87100
}
88101
}
89102

103+
private async handleListClasses(msg: { id: string }) {
104+
try {
105+
const result = (await this.client.sendRequest("modelscript/listClasses")) as {
106+
classes: { name: string; kind: string; uri: string }[];
107+
};
108+
this.panel.webview.postMessage({ type: "classListResult", id: msg.id, result });
109+
} catch (e) {
110+
this.panel.webview.postMessage({
111+
type: "classListResult",
112+
id: msg.id,
113+
result: { classes: [], error: e instanceof Error ? e.message : String(e) },
114+
});
115+
}
116+
}
117+
90118
private sendActiveFileContext() {
91-
const editor = vscode.window.activeTextEditor;
119+
// activeTextEditor is undefined when the webview panel has focus,
120+
// so fall back to visibleTextEditors to find any open Modelica file.
121+
let editor = vscode.window.activeTextEditor;
122+
if (!editor || editor.document.languageId !== "modelica") {
123+
editor = vscode.window.visibleTextEditors.find((e) => e.document.languageId === "modelica");
124+
}
92125
if (editor && editor.document.languageId === "modelica") {
93126
this.panel.webview.postMessage({
94127
type: "activeFileContext",
@@ -251,6 +284,25 @@ export class ChatPanel {
251284
40% { opacity: 1; }
252285
}
253286
287+
/* Math rendering */
288+
.math-inline {
289+
font-family: 'Cambria Math', 'Latin Modern Math', 'STIX Two Math', serif;
290+
font-style: italic;
291+
padding: 0 2px;
292+
color: var(--vscode-editor-foreground, #d4d4d4);
293+
}
294+
.math-block {
295+
font-family: 'Cambria Math', 'Latin Modern Math', 'STIX Two Math', serif;
296+
font-style: italic;
297+
display: block;
298+
text-align: center;
299+
padding: 6px 8px;
300+
margin: 4px 0;
301+
background: var(--vscode-textCodeBlock-background, #1a1a1a);
302+
border-radius: 4px;
303+
color: var(--vscode-editor-foreground, #d4d4d4);
304+
}
305+
254306
/* Input area */
255307
#input-area {
256308
display: flex;

‎packages/vscode/src/webview/chatWebview.ts‎

Lines changed: 115 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -37,24 +37,13 @@ let isGenerating = false;
3737
const conversation: ChatMessage[] = [];
3838
const pendingToolCalls = new Map<string, (result: unknown) => void>();
3939

40-
const MODEL_ID = "Qwen3-0.6B-q4f16_1-MLC";
41-
42-
const SYSTEM_PROMPT = `You are ModelScript AI, an expert Modelica language assistant integrated into the ModelScript IDE.
43-
44-
You help users with:
45-
- Writing and debugging Modelica models, packages, functions, and connectors
46-
- Understanding the Modelica Standard Library (MSL) components
47-
- Setting up simulations and interpreting results
48-
- Explaining compiler diagnostics and suggesting fixes
40+
// Workspace context (updated automatically by extension host)
41+
let activeFileName: string | null = null;
42+
let activeFileContent: string | null = null;
4943

50-
You have access to the following tools:
51-
- modelscript_flatten: Flatten a Modelica class. Use: TOOL_CALL: {"tool":"modelscript_flatten","name":"ClassName"}
52-
- modelscript_simulate: Simulate a model. Use: TOOL_CALL: {"tool":"modelscript_simulate","name":"ClassName"}
53-
- modelscript_query: Inspect a class. Use: TOOL_CALL: {"tool":"modelscript_query","name":"ClassName"}
54-
- modelscript_parse: Parse Modelica code. Use: TOOL_CALL: {"tool":"modelscript_parse","code":"model M end M;"}
44+
const MODEL_ID = "Qwen3-0.6B-q4f16_1-MLC";
5545

56-
When referencing Modelica code, use proper syntax. Be concise and helpful.
57-
Do not use <think> tags or internal reasoning blocks. Respond directly.`;
46+
const SYSTEM_PROMPT = `You are ModelScript AI, a Modelica language assistant. Answer questions about Modelica code concisely. When the user provides code context, base your answer on that specific code. Do not use <think> tags.`;
5847

5948
// ── WebLLM Engine (runs in main thread, GPU inference in internal workers) ──
6049

@@ -129,7 +118,11 @@ function addMessage(role: "user" | "assistant" | "tool", content: string): HTMLE
129118
}
130119

131120
function stripThinkTags(text: string): string {
132-
return text.replace(/<think>[\s\S]*?<\/think>\s*/g, "").trim();
121+
// Strip complete <think>...</think> blocks
122+
let cleaned = text.replace(/<think>[\s\S]*?<\/think>\s*/g, "");
123+
// Strip incomplete <think> blocks (no closing tag, model was cut off)
124+
cleaned = cleaned.replace(/<think>[\s\S]*/g, "");
125+
return cleaned.trim();
133126
}
134127

135128
function formatContent(text: string): string {
@@ -140,9 +133,42 @@ function formatContent(text: string): string {
140133
'<code style="background:var(--vscode-textCodeBlock-background,#1a1a1a);padding:1px 4px;border-radius:3px;">$1</code>',
141134
);
142135
html = html.replace(/\*\*([^*]+)\*\*/g, "<strong>$1</strong>");
136+
// Display math: $$...$$
137+
html = html.replace(/\$\$([\s\S]*?)\$\$/g, (_m, expr) => `<div class="math-block">${renderLatex(expr)}</div>`);
138+
// Inline math: $...$
139+
html = html.replace(/\$([^$\n]+)\$/g, (_m, expr) => `<span class="math-inline">${renderLatex(expr)}</span>`);
143140
return html;
144141
}
145142

143+
function renderLatex(expr: string): string {
144+
let text = expr.trim();
145+
// \frac{a}{b} → a/b
146+
text = text.replace(/\\frac\{([^}]+)\}\{([^}]+)\}/g, "($1)/($2)");
147+
// \text{...} → ...
148+
text = text.replace(/\\text\{([^}]+)\}/g, "$1");
149+
// \cdot → ·
150+
text = text.replace(/\\cdot/g, "·");
151+
// \times → ×
152+
text = text.replace(/\\times/g, "×");
153+
// \leq, \geq, \neq
154+
text = text.replace(/\\leq/g, "≤").replace(/\\geq/g, "≥").replace(/\\neq/g, "≠");
155+
// \sum, \prod, \int
156+
text = text
157+
.replace(/\\sum/g, "∑")
158+
.replace(/\\prod/g, "∏")
159+
.replace(/\\int/g, "∫");
160+
// \infty → ∞
161+
text = text.replace(/\\infty/g, "∞");
162+
// \sqrt{x} → √(x)
163+
text = text.replace(/\\sqrt\{([^}]+)\}/g, "√($1)");
164+
// \partial → ∂
165+
text = text.replace(/\\partial/g, "∂");
166+
// d(...)/dt style: keep as-is
167+
// Remove remaining backslashes from unknown commands
168+
text = text.replace(/\\([a-zA-Z]+)/g, "$1");
169+
return text;
170+
}
171+
146172
function addTypingIndicator(): HTMLElement {
147173
const div = document.createElement("div");
148174
div.className = "msg assistant typing";
@@ -172,12 +198,8 @@ window.addEventListener("message", (event) => {
172198
}
173199
break;
174200
case "activeFileContext":
175-
if (msg.content) {
176-
conversation.push({
177-
role: "system",
178-
content: `The user currently has the file "${msg.fileName}" open:\n\`\`\`modelica\n${msg.content}\n\`\`\``,
179-
});
180-
}
201+
activeFileName = msg.fileName ?? null;
202+
activeFileContent = msg.content ?? null;
181203
break;
182204
}
183205
});
@@ -193,9 +215,22 @@ async function sendMessage(): Promise<void> {
193215
inputEl.style.height = "36px";
194216
sendBtn.disabled = true;
195217

218+
// Show what the user typed, but send augmented version with context to the model
196219
addMessage("user", text);
197-
conversation.push({ role: "user", content: text });
198220

221+
// Build the augmented user message with workspace context prepended
222+
let augmentedText = "";
223+
if (activeFileName && activeFileContent) {
224+
const lines = activeFileContent.split("\n");
225+
const truncated = lines.length > 25 ? lines.slice(0, 25).join("\n") + "\n// ..." : activeFileContent;
226+
augmentedText += `Here is the code from "${activeFileName}" currently open in the editor:\n${truncated}\n\n`;
227+
}
228+
augmentedText += text;
229+
230+
conversation.push({ role: "user", content: augmentedText });
231+
console.log("[chat] context:", { activeFile: activeFileName, hasContent: !!activeFileContent });
232+
233+
// Build messages: short system prompt + conversation with augmented user messages
199234
const messages: ChatMessage[] = [{ role: "system", content: SYSTEM_PROMPT }, ...conversation];
200235

201236
const typingEl = addTypingIndicator();
@@ -206,18 +241,66 @@ async function sendMessage(): Promise<void> {
206241
const completion = await engine.chat.completions.create({
207242
messages,
208243
temperature: 0.7,
209-
max_tokens: 4096,
244+
max_tokens: 2048,
210245
stream: false,
211246
});
212247

248+
let rawText = completion.choices?.[0]?.message?.content ?? "";
249+
let finishReason = completion.choices?.[0]?.finish_reason ?? "stop";
250+
let visibleText = stripThinkTags(rawText);
251+
252+
// If the model only produced <think> content, retry once with a direct instruction
253+
if (!visibleText) {
254+
statusEl.textContent = "Retrying...";
255+
const retryMessages = [
256+
{ role: "system" as const, content: "Answer directly and concisely. No reasoning tags." },
257+
...conversation,
258+
];
259+
const retry = await engine.chat.completions.create({
260+
messages: retryMessages,
261+
temperature: 0.5,
262+
max_tokens: 2048,
263+
stream: false,
264+
});
265+
rawText = retry.choices?.[0]?.message?.content ?? "";
266+
finishReason = retry.choices?.[0]?.finish_reason ?? "stop";
267+
visibleText = stripThinkTags(rawText);
268+
}
269+
270+
// If truncated (finish_reason="length"), try one continuation with stripped content
271+
if (finishReason === "length" && visibleText) {
272+
statusEl.textContent = "Continuing...";
273+
const contMessages = [
274+
...messages,
275+
{ role: "assistant" as const, content: visibleText },
276+
{ role: "user" as const, content: "Continue." },
277+
];
278+
const cont = await engine.chat.completions.create({
279+
messages: contMessages,
280+
temperature: 0.7,
281+
max_tokens: 2048,
282+
stream: false,
283+
});
284+
const contChunk = stripThinkTags(cont.choices?.[0]?.message?.content ?? "");
285+
if (contChunk) {
286+
visibleText += " " + contChunk;
287+
rawText += cont.choices?.[0]?.message?.content ?? "";
288+
}
289+
}
290+
213291
typingEl.remove();
214-
const resultText = completion.choices?.[0]?.message?.content ?? "";
292+
statusEl.textContent = "Qwen3-0.6B ready";
215293

216-
addMessage("assistant", resultText);
217-
conversation.push({ role: "assistant", content: resultText });
294+
if (visibleText) {
295+
addMessage("assistant", visibleText);
296+
} else {
297+
addMessage("assistant", "I couldn't generate a response. Try a shorter or more specific question.");
298+
}
299+
300+
conversation.push({ role: "assistant", content: visibleText || rawText });
218301

219302
// Check for tool calls in the response
220-
const toolCallMatch = resultText.match(/TOOL_CALL:\s*(\{[\s\S]*?\})/);
303+
const toolCallMatch = visibleText.match(/TOOL_CALL:\s*(\{[\s\S]*?\})/);
221304
if (toolCallMatch) {
222305
try {
223306
const toolReq = JSON.parse(toolCallMatch[1]);
@@ -279,3 +362,6 @@ inputEl.addEventListener("input", () => {
279362
inputEl.disabled = false;
280363
sendBtn.disabled = false;
281364
inputEl.focus();
365+
366+
// Request workspace context now that the script is loaded
367+
vscode.postMessage({ type: "getActiveFileContext" });

0 commit comments

Comments
 (0)