diff --git a/apps/studio/src/lib/editor/engine/chat/context.ts b/apps/studio/src/lib/editor/engine/chat/context.ts index d492b4f60..96378a204 100644 --- a/apps/studio/src/lib/editor/engine/chat/context.ts +++ b/apps/studio/src/lib/editor/engine/chat/context.ts @@ -7,6 +7,7 @@ import { type HighlightMessageContext, type ImageMessageContext, type ProjectMessageContext, + type RelatedFileMessageContext, } from '@onlook/models/chat'; import type { DomElement } from '@onlook/models/element'; import type { ParsedError } from '@onlook/utility'; @@ -142,13 +143,70 @@ export class ChatContext { } } - getProjectContext(): ProjectMessageContext[] { + async getProjectContext(): Promise<(ProjectMessageContext | RelatedFileMessageContext)[]> { const folderPath = this.projectsManager.project?.folderPath; if (!folderPath) { return []; } - return [ + // Get selected elements + const selected = this.editorEngine.elements.selected; + if (selected.length === 0) { + return [ + { + type: MessageContextType.PROJECT, + content: '', + displayName: 'Project', + path: folderPath, + }, + ]; + } + + // Get related files from templateNodes using oid or instanceId of children + const relatedFiles = new Set(); + + // Process each selected element + for (const element of selected) { + // Get the layer node for the selected element + const layerNode = this.editorEngine.ast.mappings.getLayerNode( + element.webviewId, + element.domId, + ); + if (!layerNode || !layerNode.children) { + continue; + } + + // Process each child + for (const childId of layerNode.children) { + const childLayerNode = this.editorEngine.ast.mappings.getLayerNode( + element.webviewId, + childId, + ); + if (!childLayerNode) { + continue; + } + + // Try to get templateNode using oid + if (childLayerNode.oid) { + const templateNode = await this.editorEngine.ast.getTemplateNodeById( + childLayerNode.oid, + ); + if (templateNode && templateNode.path) { + relatedFiles.add(templateNode.path); + } + } + + // Try to get templateNode using instanceId + if (childLayerNode.instanceId) { + // For now, we'll just log that we found an instanceId + console.log(`Found child with instanceId: ${childLayerNode.instanceId}`); + // Additional logic could be added here to retrieve related files using instanceId + } + } + } + + // Create project context with related files + const projectContext: (ProjectMessageContext | RelatedFileMessageContext)[] = [ { type: MessageContextType.PROJECT, content: '', @@ -156,6 +214,23 @@ export class ChatContext { path: folderPath, }, ]; + + // Add related files to project context + for (const filePath of relatedFiles) { + const fileContent = await this.editorEngine.code.getFileContent(filePath, true); + if (fileContent === null) { + continue; + } + + projectContext.push({ + type: MessageContextType.RELATED_FILE, + content: fileContent, + displayName: `Related: ${filePath}`, + path: filePath, + }); + } + + return projectContext; } getMessageContext(errors: ParsedError[]): ErrorMessageContext[] { diff --git a/apps/studio/src/lib/editor/engine/element/index.ts b/apps/studio/src/lib/editor/engine/element/index.ts index ae97a79b1..176d438db 100644 --- a/apps/studio/src/lib/editor/engine/element/index.ts +++ b/apps/studio/src/lib/editor/engine/element/index.ts @@ -130,6 +130,17 @@ export class ElementManager { this.selectedElements = []; } + getElementByDomId(domId: string): DomElement | null { + // Check all webviews for the element + for (const webviewId of this.editorEngine.webviews.getWebviewIds()) { + const layerNode = this.editorEngine.ast.mappings.getLayerNode(webviewId, domId); + if (layerNode) { + return layerNode; + } + } + return null; + } + async delete() { const selected = this.selected; if (selected.length === 0) { diff --git a/packages/models/src/chat/message/context.ts b/packages/models/src/chat/message/context.ts index 868bda821..cae76a49d 100644 --- a/packages/models/src/chat/message/context.ts +++ b/packages/models/src/chat/message/context.ts @@ -4,6 +4,7 @@ export enum MessageContextType { IMAGE = 'image', ERROR = 'error', PROJECT = 'project', + RELATED_FILE = 'relatedFile', } type BaseMessageContext = { @@ -38,9 +39,15 @@ export type ProjectMessageContext = BaseMessageContext & { path: string; }; +export type RelatedFileMessageContext = BaseMessageContext & { + type: MessageContextType.RELATED_FILE; + path: string; +}; + export type ChatMessageContext = | FileMessageContext | HighlightMessageContext | ImageMessageContext | ErrorMessageContext - | ProjectMessageContext; + | ProjectMessageContext + | RelatedFileMessageContext;