diff --git a/docs/query-generation-data-flow.md b/docs/query-generation-data-flow.md new file mode 100644 index 000000000..2d7d8a7d7 --- /dev/null +++ b/docs/query-generation-data-flow.md @@ -0,0 +1,231 @@ +# Query Generation Data Flow - Privacy Review Documentation + +## Overview + +This document describes the data flow for the MongoDB query generation feature, with emphasis on customer data handling and privacy considerations. The query generation feature uses GitHub Copilot's language models to convert natural language descriptions into MongoDB queries. + +## Current Implementation (v1.0) + +### Data Flow Diagram + +``` +┌─────────────────┐ +│ User Input │ +│ (Natural Lang) │ +└────────┬────────┘ + │ + v +┌─────────────────────────────────────────────────────────┐ +│ Query Generation Context │ +│ - Database name │ +│ - Collection name(s) │ +│ - Target query type (Find/Aggregation) │ +│ - Natural language query description │ +└────────┬────────────────────────────────────────────────┘ + │ + v +┌─────────────────────────────────────────────────────────┐ +│ Schema Inference Process │ +│ 1. Fetch sample documents from customer database │ +│ 2. Analyze document structure locally │ +│ 3. Generate schema definition (field types only) │ +│ 4. DISCARD original documents │ +└────────┬────────────────────────────────────────────────┘ + │ + v +┌─────────────────────────────────────────────────────────┐ +│ Data Sent to LLM (GitHub Copilot) │ +│ ✓ Database name (metadata) │ +│ ✓ Collection name(s) (metadata) │ +│ ✓ Schema structure (field names and types only) │ +│ ✓ Natural language query from user │ +│ ✗ NO actual customer data values │ +│ ✗ NO sample documents │ +└────────┬────────────────────────────────────────────────┘ + │ + v +┌─────────────────────────────────────────────────────────┐ +│ GitHub Copilot LLM Processing │ +│ - Processes schema structure │ +│ - Generates MongoDB query syntax │ +└────────┬────────────────────────────────────────────────┘ + │ + v +┌─────────────────────────────────────────────────────────┐ +│ Response │ +│ - Generated MongoDB query (JSON) │ +│ - Explanation of the query logic │ +└────────┬────────────────────────────────────────────────┘ + │ + v +┌─────────────────┐ +│ Display to │ +│ User │ +└─────────────────┘ +``` + +### Customer Data Categories + +#### Data Collected Locally (Never Sent to LLM) + +1. **Sample Documents**: + - Fetched: 3-10 documents per collection + - Purpose: Schema inference only + - Lifecycle: Used for schema analysis, then immediately discarded + - Storage: Temporary in-memory only, never persisted + +2. **Actual Data Values**: + - Type: Any customer data content (strings, numbers, objects, etc.) + - Handling: Never extracted, never sent to LLM + - Example: If a document has `{"name": "John Doe", "age": 30}`, these values are never sent + +#### Data Sent to LLM + +1. **Metadata**: + - Database name (e.g., "customerDB") + - Collection name(s) (e.g., "users", "orders") + - Target query type ("Find" or "Aggregation") + +2. **Schema Structure**: + - Field names (e.g., "name", "age", "email") + - Field types (e.g., "string", "number", "object") + - Nested structure (e.g., "address.city" is a string) + - **NO actual values from customer documents** + +3. **User Input**: + - Natural language query description provided by the user + - Example: "Find all users who are over 25 years old" + +### Schema Inference Example + +#### Customer's Actual Document (Never Sent): +```json +{ + "_id": "507f1f77bcf86cd799439011", + "name": "John Doe", + "email": "john.doe@example.com", + "age": 30, + "address": { + "city": "Seattle", + "state": "WA", + "zipCode": "98101" + }, + "orders": [ + {"orderId": "ORD-001", "total": 99.99} + ] +} +``` + +#### Schema Definition Sent to LLM: +```json +{ + "collectionName": "users", + "fields": { + "_id": "string", + "name": "string", + "email": "string", + "age": "number", + "address": { + "city": "string", + "state": "string", + "zipCode": "string" + }, + "orders": [ + { + "orderId": "string", + "total": "number" + } + ] + } +} +``` + +**Key Privacy Point**: Only the structure (field names and types) is sent. Actual values like "John Doe", "john.doe@example.com", "Seattle", etc., are never included in the LLM request. + +### Code Reference + +The schema inference is implemented in `src/utils/schemaInference.ts`: + +```typescript +export function generateSchemaDefinition( + documents: Array, + collectionName?: string, +): SchemaDefinition { + // Processes documents to extract ONLY field names and types + // Returns structure without any actual data values +} +``` + +The query generation call in `src/commands/llmEnhancedCommands/queryGenerationCommands.ts`: + +```typescript +// Sample documents are fetched +const sampleDocs = await client.getSampleDocuments( + queryContext.databaseName, + queryContext.collectionName, + 10 +); + +// Schema is extracted (structure only) +const schema = generateSchemaDefinition(sampleDocs, queryContext.collectionName); + +// Original documents are discarded after this point +// Only schema structure is used in prompt template +``` + +## Proposed Future Enhancement (v2.0) - Under Privacy Review + +### Overview of Proposed Change + +To enable query modification features, we are considering allowing users to provide their existing MongoDB queries for modification or optimization. + +### Additional Data Flow (Proposed) + +``` +┌─────────────────────────────────────┐ +│ User Input (New) │ +│ - Natural language modification │ +│ request │ +│ - Existing MongoDB query (CUSTOMER │ +│ CREATED, may contain literals) │ +└────────┬────────────────────────────┘ + │ + v +┌─────────────────────────────────────┐ +│ Data Sent to LLM (Additional) │ +│ ✓ User's original query structure │ +│ ⚠ May contain customer-specified │ +│ literals/values in query filters │ +└─────────────────────────────────────┘ +``` + +### Privacy Concerns with Proposed Enhancement + +#### New Customer Data Being Sent to LLM: + +1. **User's Original Query**: + - Content: MongoDB query syntax provided by the user + - Risk: May contain literal values used in filters + - Example: + ```javascript + // User's query may contain: + db.users.find({ + "email": "specific@customer.com", // ⚠ Customer email + "accountId": "ACCT-12345" // ⚠ Customer account ID + }) + ``` + +2. **Embedded Literals**: + - Query filters often contain specific values + - These values could be sensitive customer data + - Examples: email addresses, account IDs, names, dates, amounts + +#### Privacy Risk Assessment: + +| Data Type | Current (v1.0) | Proposed (v2.0) | Risk Level | +|-----------|----------------|-----------------|------------| +| Sample document values | ✗ Never sent | ✗ Never sent | None | +| Schema structure | ✓ Sent | ✓ Sent | Low (metadata only) | +| Database/collection names | ✓ Sent | ✓ Sent | Low (metadata) | +| User's natural language input | ✓ Sent | ✓ Sent | Low-Medium (user provided) | +| Query literals/filters | ✗ Not applicable | ⚠ **Would be sent** | **Medium-High** | diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 8c970c42f..8219978e8 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -87,6 +87,7 @@ "Change page size": "Change page size", "Changelog": "Changelog", "Check document syntax": "Check document syntax", + "Checking GitHub Copilot availability...": "Checking GitHub Copilot availability...", "Choose a cluster…": "Choose a cluster…", "Choose a RU cluster…": "Choose a RU cluster…", "Choose a Subscription…": "Choose a Subscription…", @@ -99,11 +100,13 @@ "Click here to update credentials": "Click here to update credentials", "Click to view resource": "Click to view resource", "Close the account management wizard": "Close the account management wizard", + "Cluster metadata not initialized. Client may not be properly connected.": "Cluster metadata not initialized. Client may not be properly connected.", "Cluster support unknown $(info)": "Cluster support unknown $(info)", "Collection name cannot begin with the system. prefix (Reserved for internal use).": "Collection name cannot begin with the system. prefix (Reserved for internal use).", "Collection name cannot contain .system.": "Collection name cannot contain .system.", "Collection name cannot contain the $.": "Collection name cannot contain the $.", "Collection name cannot contain the null character.": "Collection name cannot contain the null character.", + "Collection name is required for single-collection query generation": "Collection name is required for single-collection query generation", "Collection name is required.": "Collection name is required.", "Collection names should begin with an underscore or a letter character.": "Collection names should begin with an underscore or a letter character.", "Configure Azure Discovery Filters": "Configure Azure Discovery Filters", @@ -115,7 +118,9 @@ "Configuring tenant filtering…": "Configuring tenant filtering…", "Connect to a database": "Connect to a database", "Connected to \"{name}\"": "Connected to \"{name}\"", + "Connected to cluster [{durationMs}ms]": "Connected to cluster [{durationMs}ms]", "Connected to the cluster \"{cluster}\".": "Connected to the cluster \"{cluster}\".", + "Connecting to cluster: {clusterId}": "Connecting to cluster: {clusterId}", "Connecting to the cluster as \"{username}\"…": "Connecting to the cluster as \"{username}\"…", "Connecting to the cluster using Entra ID…": "Connecting to the cluster using Entra ID…", "Connection String": "Connection String", @@ -156,6 +161,7 @@ "Deleting...": "Deleting...", "Disable TLS/SSL (Not recommended)": "Disable TLS/SSL (Not recommended)", "Disable TLS/SSL checks when connecting.": "Disable TLS/SSL checks when connecting.", + "Discovering collections in database: {databaseName}": "Discovering collections in database: {databaseName}", "Do not rely on case to distinguish between databases. For example, you cannot use two databases with names like, salesData and SalesData.": "Do not rely on case to distinguish between databases. For example, you cannot use two databases with names like, salesData and SalesData.", "Do not save credentials.": "Do not save credentials.", "Document must be an object.": "Document must be an object.", @@ -235,10 +241,16 @@ "Failed to extract cluster credentials from the selected node.": "Failed to extract cluster credentials from the selected node.", "Failed to extract the connection string from the selected account.": "Failed to extract the connection string from the selected account.", "Failed to find commandId on generic tree item.": "Failed to find commandId on generic tree item.", + "Failed to gather schema information: {message}": "Failed to gather schema information: {message}", "Failed to get public IP": "Failed to get public IP", "Failed to initialize Azure management clients": "Failed to initialize Azure management clients", + "Failed to load custom prompt template from {path}: {error}. Using built-in template.": "Failed to load custom prompt template from {path}: {error}. Using built-in template.", + "Failed to load template file for {type}: {error}": "Failed to load template file for {type}: {error}", "Failed to obtain Entra ID token.": "Failed to obtain Entra ID token.", + "Failed to parse generated query. Query generation provided an invalid response.": "Failed to parse generated query. Query generation provided an invalid response.", + "Failed to parse language model response": "Failed to parse language model response", "Failed to parse secrets for key {0}:": "Failed to parse secrets for key {0}:", + "Failed to parse the response from the language model. LLM output:\n{output}": "Failed to parse the response from the language model. LLM output:\n{output}", "Failed to process URI: {0}": "Failed to process URI: {0}", "Failed to rename the connection.": "Failed to rename the connection.", "Failed to retrieve Azure accounts: {0}": "Failed to retrieve Azure accounts: {0}", @@ -247,10 +259,18 @@ "Failed to store secrets for key {0}:": "Failed to store secrets for key {0}:", "Failed to update the connection.": "Failed to update the connection.", "Failed with code \"{0}\".": "Failed with code \"{0}\".", + "Filling prompt template...": "Filling prompt template...", "Find Query": "Find Query", "Finished importing": "Finished importing", + "Found {count} collections [{durationMs}ms]": "Found {count} collections [{durationMs}ms]", + "Gathering schema information...": "Gathering schema information...", "Generate": "Generate", "Generate query with AI": "Generate query with AI", + "Generating schema from {count} sample documents for collection: {collectionName}": "Generating schema from {count} sample documents for collection: {collectionName}", + "Generating schema from sample documents...": "Generating schema from sample documents...", + "GitHub Copilot is available": "GitHub Copilot is available", + "GitHub Copilot is not available": "GitHub Copilot is not available", + "GitHub Copilot is not available. Please install the GitHub Copilot extension and ensure you have an active subscription.": "GitHub Copilot is not available. Please install the GitHub Copilot extension and ensure you have an active subscription.", "Go back.": "Go back.", "Go to first page": "Go to first page", "Go to next page": "Go to next page", @@ -346,10 +366,12 @@ "No properties found in the schema at path \"{0}\"": "No properties found in the schema at path \"{0}\"", "No public connectivity": "No public connectivity", "No result returned from the MongoDB shell.": "No result returned from the MongoDB shell.", + "No sample documents found for collection: {collectionName}": "No sample documents found for collection: {collectionName}", "No scope was provided for the role assignment.": "No scope was provided for the role assignment.", "No session found for id {sessionId}": "No session found for id {sessionId}", "No subscriptions found": "No subscriptions found", "No subscriptions found for the selected tenants. Please adjust your tenant selection or check your Azure permissions.": "No subscriptions found for the selected tenants. Please adjust your tenant selection or check your Azure permissions.", + "No suitable language model found. Please ensure GitHub Copilot is installed and you have an active subscription.": "No suitable language model found. Please ensure GitHub Copilot is installed and you have an active subscription.", "No tenants found. Please try signing in again or check your Azure permissions.": "No tenants found. Please try signing in again or check your Azure permissions.", "No tenants selected. Azure discovery will be filtered to exclude all tenant results.": "No tenants selected. Azure discovery will be filtered to exclude all tenant results.", "Not connected to any MongoDB database.": "Not connected to any MongoDB database.", @@ -360,6 +382,7 @@ "Open installation page": "Open installation page", "Opening DocumentDB connection…": "Opening DocumentDB connection…", "Operation cancelled.": "Operation cancelled.", + "Parsing language model response...": "Parsing language model response...", "Password for {username_at_resource}": "Password for {username_at_resource}", "Pick \"{number}\" to confirm and continue.": "Pick \"{number}\" to confirm and continue.", "Please authenticate first by expanding the tree item of the selected cluster.": "Please authenticate first by expanding the tree item of the selected cluster.", @@ -378,8 +401,15 @@ "Procedure not found: {name}": "Procedure not found: {name}", "Process exited: \"{command}\"": "Process exited: \"{command}\"", "Project": "Project", + "Prompt template filled successfully": "Prompt template filled successfully", "Provide Feedback": "Provide Feedback", "Provider \"{0}\" does not have resource type \"{1}\".": "Provider \"{0}\" does not have resource type \"{1}\".", + "Query generation completed successfully": "Query generation completed successfully", + "Query generation failed": "Query generation failed", + "Query generation failed with the error: {0}": "Query generation failed with the error: {0}", + "Query generation is using model \"{actualModel}\" instead of preferred \"{preferredModel}\". Results may vary.": "Query generation is using model \"{actualModel}\" instead of preferred \"{preferredModel}\". Results may vary.", + "Query generation started: type={type}, targetQueryType={targetQueryType}": "Query generation started: type={type}, targetQueryType={targetQueryType}", + "Received response from model: {modelUsed} [{durationMs}ms]": "Received response from model: {modelUsed} [{durationMs}ms]", "Refresh": "Refresh", "Refresh current view": "Refresh current view", "Refreshing Azure discovery tree…": "Refreshing Azure discovery tree…", @@ -391,12 +421,16 @@ "Report a Bug": "Report a Bug", "Report an issue": "Report an issue", "Resource group \"{0}\" already exists in subscription \"{1}\".": "Resource group \"{0}\" already exists in subscription \"{1}\".", + "Retrieved {count} sample documents [{durationMs}ms]": "Retrieved {count} sample documents [{durationMs}ms]", "Return to the account list": "Return to the account list", "Reusing active connection for \"{cluster}\".": "Reusing active connection for \"{cluster}\".", "Revisit connection details and try again.": "Revisit connection details and try again.", "Role assignment \"{0}\" created for the {2} resource \"{1}\".": "Role assignment \"{0}\" created for the {2} resource \"{1}\".", "Role Assignment {0} created for {1}": "Role Assignment {0} created for {1}", "Role Assignment {0} failed for {1}": "Role Assignment {0} failed for {1}", + "Sampled {count} documents from collection: {collectionName} [{durationMs}ms]": "Sampled {count} documents from collection: {collectionName} [{durationMs}ms]", + "Sampling documents from collection: {collectionName}": "Sampling documents from collection: {collectionName}", + "Sampling documents from collection: {databaseName}.{collectionName}": "Sampling documents from collection: {databaseName}.{collectionName}", "Save": "Save", "Save credentials for future connections.": "Save credentials for future connections.", "Save credentials for future use?": "Save credentials for future use?", @@ -404,6 +438,9 @@ "Save to the database": "Save to the database", "Saving \"{path}\" will update the entity \"{name}\" to the cloud.": "Saving \"{path}\" will update the entity \"{name}\" to the cloud.", "Saving credentials for \"{clusterName}\"…": "Saving credentials for \"{clusterName}\"…", + "Schema gathering completed [{durationMs}ms]": "Schema gathering completed [{durationMs}ms]", + "Schema gathering failed: {message}": "Schema gathering failed: {message}", + "Schema generation complete for {count} collection(s)": "Schema generation complete for {count} collection(s)", "Select {0}": "Select {0}", "Select {mongoExecutableFileName}": "Select {mongoExecutableFileName}", "Select a location for new resources.": "Select a location for new resources.", @@ -421,6 +458,7 @@ "Select the local connection type…": "Select the local connection type…", "Selected subscriptions: {0}": "Selected subscriptions: {0}", "Selected tenants: {0}": "Selected tenants: {0}", + "Sending request to GitHub Copilot (preferred model: {preferredModel})...": "Sending request to GitHub Copilot (preferred model: {preferredModel})...", "Service Discovery": "Service Discovery", "Sign in to Azure to continue…": "Sign in to Azure to continue…", "Sign in to Azure...": "Sign in to Azure...", @@ -454,6 +492,8 @@ "Tag can only contain alphanumeric characters, underscores, periods, and hyphens.": "Tag can only contain alphanumeric characters, underscores, periods, and hyphens.", "Tag cannot be empty.": "Tag cannot be empty.", "Tag cannot be longer than 256 characters.": "Tag cannot be longer than 256 characters.", + "Template file is empty: {path}": "Template file is empty: {path}", + "Template file not found: {path}": "Template file not found: {path}", "Tenant ID cannot be empty": "Tenant ID cannot be empty", "Tenant ID: {0}": "Tenant ID: {0}", "Tenant Name: {0}": "Tenant Name: {0}", @@ -516,6 +556,7 @@ "Unexpected status code: {0}": "Unexpected status code: {0}", "Unknown error": "Unknown error", "Unknown Error": "Unknown Error", + "Unknown query generation type: {type}": "Unknown query generation type: {type}", "Unknown tenant": "Unknown tenant", "Unrecognized node type encountered. Could not parse {constructorCall} as part of {functionCall}": "Unrecognized node type encountered. Could not parse {constructorCall} as part of {functionCall}", "Unrecognized node type encountered. We could not parse {text}": "Unrecognized node type encountered. We could not parse {text}", @@ -525,6 +566,7 @@ "Unsupported authentication method: {0}": "Unsupported authentication method: {0}", "Unsupported authentication method.": "Unsupported authentication method.", "Unsupported emulator type: \"{emulatorType}\"": "Unsupported emulator type: \"{emulatorType}\"", + "Unsupported query type: {queryType}": "Unsupported query type: {queryType}", "Unsupported resource: {0}": "Unsupported resource: {0}", "Unsupported view for an authentication retry.": "Unsupported view for an authentication retry.", "Up": "Up", @@ -538,6 +580,7 @@ "Username and Password": "Username and Password", "Username cannot be empty": "Username cannot be empty", "Username for {resource}": "Username for {resource}", + "Using custom prompt template for {type} query generation: {path}": "Using custom prompt template for {type} query generation: {path}", "Using existing resource group \"{0}\".": "Using existing resource group \"{0}\".", "Using the table navigation, you can explore deeper levels or move back and forth between them.": "Using the table navigation, you can explore deeper levels or move back and forth between them.", "Validate": "Validate", diff --git a/package.json b/package.json index ce2b61c10..d60323e48 100644 --- a/package.json +++ b/package.json @@ -900,6 +900,46 @@ "type": "number", "description": "The batch size to be used when querying working with the shell.", "default": 50 + }, + "documentDB.aiAssistant.findQueryPromptPath": { + "type": [ + "string", + "null" + ], + "description": "Path to a custom prompt template file for find query optimization. Leave empty to use the built-in template.", + "default": null + }, + "documentDB.aiAssistant.aggregateQueryPromptPath": { + "type": [ + "string", + "null" + ], + "description": "Path to a custom prompt template file for aggregate query optimization. Leave empty to use the built-in template.", + "default": null + }, + "documentDB.aiAssistant.countQueryPromptPath": { + "type": [ + "string", + "null" + ], + "description": "Path to a custom prompt template file for count query optimization. Leave empty to use the built-in template.", + "default": null + }, + "documentDB.aiAssistant.crossCollectionQueryPromptPath": { + "type": [ + "string", + "null" + ], + "description": "Path to a custom prompt template file for cross-collection query generation. Leave empty to use the built-in template.", + "default": null + }, + "documentDB.aiAssistant.singleCollectionQueryPromptPath": { + "type": [ + "string", + "null" + ], + "description": "Path to a custom prompt template file for single collection query generation. Leave empty to use the built-in template.", + "default": null } } } diff --git a/src/commands/llmEnhancedCommands/promptTemplates.ts b/src/commands/llmEnhancedCommands/promptTemplates.ts new file mode 100644 index 000000000..86d703880 --- /dev/null +++ b/src/commands/llmEnhancedCommands/promptTemplates.ts @@ -0,0 +1,518 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { l10n } from 'vscode'; + +/** + * Preferred language model for index optimization + */ +export const PREFERRED_MODEL = 'gpt-5'; + +/** + * Fallback models to use if the preferred model is not available + */ +export const FALLBACK_MODELS = ['gpt-4o', 'gpt-4o-mini']; + +/** + * Embedded prompt templates for query optimization + * These templates are compiled into the extension bundle at build time + */ + +export const FIND_QUERY_PROMPT_TEMPLATE = ` +You are an expert MongoDB assistant to provide index suggestions for the following find query: +- **Query**: {query} +The query is executed against a MongoDB collection with the following details: +## Cluster Information +- **Is_Azure_Cluster**: {isAzureCluster} +- **Azure_Cluster_Type**: {AzureClusterType} +## Collection Information +- **Collection_Stats**: {collectionStats} +## Index Information of Current Collection +- **Indexes_Stats**: {indexStats} +## Query Execution Stats +- **Execution_Stats**: {executionStats} +Follow these strict instructions (must obey): +1. **Single JSON output only** — your response MUST be a single valid JSON object and **nothing else** (no surrounding text, no code fences, no explanation). +2. **Do not hallucinate** — only use facts present in the sections Query, Collection_Stats, Indexes_Stats, Execution_Stats. If a required metric is absent, set the corresponding field to \`null\` in \`metadata\`. +3. **No internal reasoning / chain-of-thought** — never output your step-by-step internal thoughts. Give concise, evidence-based conclusions only. +4. **Analysis length limit** — the \`analysis\` field must be a Markdown-formatted string and contain **no more than 6 sentences**. Be concise. +5. **Runnable shell commands** — any index changes you recommend must be provided as **mongosh/mongo shell** commands (runnable). Use \`db.getCollection("{collectionName}")\` to reference the collection (replace \`{collectionName}\` with the actual name from \`collectionStats\`). +6. **Justify every index command** — each \`create\`/\`drop\` recommendation must include a one-sentence justification that references concrete fields/metrics from \`executionStats\` or \`indexStats\`. +7. **Prefer minimal, safe changes** — prefer a single, high-impact index over many small ones; avoid suggesting drops unless the benefit is clear and justified. +8. **Include priority** — each suggested improvement must include a \`priority\` (\`high\`/\`medium\`/\`low\`) so an engineer can triage. +9. **Be explicit about risks** — if a suggested index could increase write cost or large index size, include that as a short risk note in the improvement. +10. **Verification output** — the \`verification\` field must be a **Markdown string** (not an array). It should include one or more \`\`\`javascript code blocks\`\`\` containing **valid mongosh commands** to verify index performance or collection stats. Each command must be copy-paste runnable in mongosh (e.g. \`db.getCollection("{collectionName}").find(...).hint(...).explain("executionStats")\`). +11. **Do not change input objects** — echo input objects only under \`metadata\`; do not mutate \`{collectionStats}\`, \`{indexStats}\`, or \`{executionStats}\`—just include them as-is (and add computed helper fields if needed). +12. **Drop indexes with index Name** — when you drop an index, use the **index name** to reference it, not the field name. +13. **If no change recommended** — return an empty \`improvements\` array and still include a short Markdown \`verification\` section to confirm the current plan. +Thinking / analysis tips (useful signals to form recommendations; don't output these tips themselves): +- Check **which index(es)** the winning plan used (or whether a COLLSCAN occurred) and whether \`totalKeysExamined\` is much smaller than \`totalDocsExamined\` (indicates good index filtering vs heavy document fetch). +- Look for **equality predicates vs range predicates**: equality fields should be placed before range fields in compound indexes for best selectivity. +- Match **sort order** to index order to avoid blocking in-memory sorts — if query sorts on \`a:1, b:-1\` prefer an index with the same field order/direction. +- Consider **projection coverage**: if the projection only uses indexed fields, a covered (index-only) plan is possible — prefer indexes that cover both filters and projected fields. +- Beware **multikey / array** fields and sparse data — multikey fields affect index ordering and whether index-only is achievable. +- For \`$or\` branches, check if index intersection or separate indexes per branch is better; prefer a single compound index when branches share the same leading predicates. +- Consider **index size and write amplification** — if proposed index keys are high-cardinality but cover few queries, prefer partial or sparse indexes or a more selective prefix. +- For aggregation pipelines, identify whether early \`$match\`/\`$sort\` stages can benefit from indexes (match-before-project, sort after match). +- Avoid recommending duplicate or superseded indexes — check \`indexStats\` names and key patterns first. +- If the input query contains \`sort\`, \`projection\`, or aggregation stages, account for them when recommending index key order and coverage. +- If you identify indexes related to the query that have **not been accessed for a long time** or **are not selective**, consider recommending **dropping** them to reduce write and storage overhead. +- If you identify query is on a **small collection** (e.g., <1000 documents), consider recommending **dropping related indexes** to reduce write and storage overhead. +- If the **Azure_Cluster_Type** is "vCore" and a **composite index** is being created, include in \`indexOptions\` the setting: \`"storageEngine": { "enableOrderedIndex": true }\`. +Output JSON schema (required shape; **adhere exactly**): +\`\`\` +{ + "metadata": { + "collectionName": "", + "collectionStats": { ... }, + "indexStats": [ ... ], + "executionStats": { ... }, + "derived": { + "totalKeysExamined": , + "totalDocsExamined": , + "keysToDocsRatio": , + "usedIndex": "" + } + }, + "analysis": "", + "improvements": [ + { + "action": "create" | "drop" | "none" | "modify", + "indexSpec": { "": 1|-1, ... }, + "indexOptions": { }, + "mongoShell": "db.getCollection(\\"{collectionName}\\").createIndex({...}, {...})" , + "justification": "", + "priority": "high" | "medium" | "low", + "risks": "" + } + ], + "verification": "" +} +\`\`\` +Additional rules for the JSON: +- \`metadata.collectionName\` must be filled from \`{collectionStats.ns}\` or a suitable field; if not available set to \`null\`. +- \`derived.totalKeysExamined\`, \`derived.totalDocsExamined\`, and \`derived.keysToDocsRatio\` should be filled from \`executionStats\` if present, otherwise \`null\`. \`keysToDocsRatio\` = \`totalKeysExamined / max(1, totalDocsExamined)\`. +- \`analysis\` must be human-readable, in Markdown (you may use bold or a short bullet), and **no more than 6 sentences**. +- \`mongoShell\` commands must **only** use double quotes and valid JS object notation. +- \`verification\` must be human-readable, in Markdown. It should include one or more \`\`\`javascript code blocks\`\`\` containing valid mongosh commands. Each code block should be concise and executable as-is in mongosh. +`; + +export const AGGREGATE_QUERY_PROMPT_TEMPLATE = ` +You are an expert MongoDB assistant to provide index suggestions for the following aggregation pipeline: +- **Pipeline**: {pipeline} +The pipeline is executed against a MongoDB collection with the following details: +## Cluster Information +- **Is_Azure_Cluster**: {isAzureCluster} +- **Azure_Cluster_Type**: {AzureClusterType} +## Collection Information +- **Collection_Stats**: {collectionStats} +## Index Information of Current Collection +- **Indexes_Stats**: {indexStats} +## Query Execution Stats +- **Execution_Stats**: {executionStats} +## Cluster Information +- **Cluster_Type**: {clusterType} // e.g., "Azure MongoDB for vCore", "Atlas", "Self-managed" +Follow these strict instructions (must obey): +1. **Single JSON output only** — your response MUST be a single valid JSON object and **nothing else** (no surrounding text, no code fences, no explanation). +2. **Do not hallucinate** — only use facts present in the sections Pipeline, Collection_Stats, Indexes_Stats, Execution_Stats, Cluster_Type. If a required metric is absent, set the corresponding field to \`null\` in \`metadata\`. +3. **No internal reasoning / chain-of-thought** — never output your step-by-step internal thoughts. Give concise, evidence-based conclusions only. +4. **Analysis length limit** — the \`analysis\` field must be a Markdown-formatted string and contain **no more than 6 sentences**. Be concise. +5. **Runnable shell commands** — any index changes you recommend must be provided as **mongosh/mongo shell** commands (runnable). Use \`db.getCollection("{collectionName}")\` to reference the collection (replace \`{collectionName}\` with the actual name from \`collectionStats\`). +6. **Justify every index command** — each \`create\`/\`drop\` recommendation must include a one-sentence justification that references concrete fields/metrics from \`executionStats\` or \`indexStats\`. +7. **Prefer minimal, safe changes** — prefer a single, high-impact index over many small ones; avoid suggesting drops unless the benefit is clear and justified. +8. **Include priority** — each suggested improvement must include a \`priority\` (\`high\`/\`medium\`/\`low\`) so an engineer can triage. +9. **Be explicit about risks** — if a suggested index could increase write cost or large index size, include that as a short risk note in the improvement. +10. **Verification output** — the \`verification\` field must be a **Markdown string** (not an array). It should include one or more \`\`\`javascript code blocks\`\`\` containing **valid mongosh commands** to verify index performance or collection stats. Each command must be copy-paste runnable in mongosh (e.g. \`db.getCollection("{collectionName}").find(...).hint(...).explain("executionStats")\`). +11. **Do not change input objects** — echo input objects only under \`metadata\`; do not mutate \`{collectionStats}\`, \`{indexStats}\`, or \`{executionStats}\`—just include them as-is (and add computed helper fields if needed). +12. **If no change recommended** — return an empty \`improvements\` array and still include a short Markdown \`verification\` section to confirm the current plan. +Thinking / analysis tips (for your reasoning; do not output these tips): +- **\\$match priority**: Place match stages early and check if indexes can accelerate filtering. +- **\\$sort optimization**: Match sort order to index order to avoid blocking in-memory sorts. +- **\\$group / \\$project coverage**: Check if fields used in group or project stages are covered by indexes for potential index-only plans. +- **\\$lookup / \\$unwind**: Evaluate whether join or array-unwind stages can benefit from supporting indexes. +- **Multi-branch match**: For \\$or or \\$in conditions, consider compound indexes or index intersection. +- **Multikey / sparse / partial indexes**: Ensure indexes on array or sparse fields still support coverage without excessive size or write amplification. +- **Index size and write cost**: Avoid high-cardinality indexes that rarely match queries; prefer selective prefixes or partial indexes. +- **Projection coverage**: If all projected fields are indexed, prioritize index-only scan opportunities. +- If you identify indexes related to the query that have **not been accessed for a long time** or **are not selective**, consider recommending **dropping** them to reduce write and storage overhead. +- If you identify query is on a **small collection** (e.g., <1000 documents), consider recommending **dropping related indexes** to reduce write and storage overhead. +- If the **Azure_Cluster_Type** is "vCore" and a **composite index** is being created, include in \`indexOptions\` the setting: \`"storageEngine": { "enableOrderedIndex": true }\`. +- **Vector recall rule** — If the **Azure_Cluster_Type** is "vCore" and uses a cosmosSearch with index has \`"kind": "vector-ivf"\`, but the collection contains many documents (over 10k) or the vector dimensionality is high, recommend replacing it with a \`vector-hnsw\` index for better recall and retrieval quality. The recommended creation command format is: + { + "createIndexes": "", + "indexes": [ + { + "name": "", + "key": { + "": "cosmosSearch" + }, + "cosmosSearchOptions": { + "kind": "vector-hnsw", + "m": , + "efConstruction": , + "similarity": "", + "dimensions": + } + } + ] + } +Output JSON schema (required shape; adhere exactly): +\`\`\` +{ + "metadata": { + "collectionName": "", + "collectionStats": { ... }, + "indexStats": [ ... ], + "executionStats": { ... }, + "derived": { + "totalKeysExamined": , + "totalDocsExamined": , + "keysToDocsRatio": , + "usedIndex": "" + } + }, + "analysis": "", + "improvements": [ + { + "action": "create" | "drop" | "none" | "modify", + "indexSpec": { "": 1|-1, ... }, + "indexOptions": { }, + "mongoShell": "db.getCollection(\\"{collectionName}\\").createIndex({...}, {...})" , + "justification": "", + "priority": "high" | "medium" | "low", + "risks": "" + } + ], + "verification": "" +} +\`\`\` +Additional rules for the JSON: +- \`metadata.collectionName\` must be filled from \`{collectionStats.ns}\` or a suitable field; if not available set to \`null\`. +- \`derived.totalKeysExamined\`, \`derived.totalDocsExamined\`, and \`derived.keysToDocsRatio\` should be filled from \`executionStats\` if present, otherwise \`null\`. \`keysToDocsRatio\` = \`totalKeysExamined / max(1, totalDocsExamined)\`. +- \`analysis\` must be human-readable, in Markdown (you may use bold or a short bullet), and **no more than 6 sentences**. +- \`mongoShell\` commands must **only** use double quotes and valid JS object notation. +- \`verification\` must be human-readable, in Markdown. It should include one or more \`\`\`javascript code blocks\`\`\` containing valid mongosh commands. Each code block should be concise and executable as-is in mongosh. +`; + +export const COUNT_QUERY_PROMPT_TEMPLATE = ` +You are an expert MongoDB assistant to provide index suggestions for the following count query: +- **Query**: {query} +The query is executed against a MongoDB collection with the following details: +## Cluster Information +- **Is_Azure_Cluster**: {isAzureCluster} +- **Azure_Cluster_Type**: {AzureClusterType} +## Collection Information +- **Collection_Stats**: {collectionStats} +## Index Information of Current Collection +- **Indexes_Stats**: {indexStats} +## Query Execution Stats +- **Execution_Stats**: {executionStats} +## Cluster Information +- **Cluster_Type**: {clusterType} // e.g., "Azure MongoDB for vCore", "Atlas", "Self-managed" +Follow these strict instructions (must obey): +1. **Single JSON output only** — your response MUST be a single valid JSON object and **nothing else** (no surrounding text, no code fences, no explanation). +2. **Do not hallucinate** — only use facts present in the sections Query, Collection_Stats, Indexes_Stats, Execution_Stats, Cluster_Type. If a required metric is absent, set the corresponding field to \`null\` in \`metadata\`. +3. **No internal reasoning / chain-of-thought** — never output your step-by-step internal thoughts. Give concise, evidence-based conclusions only. +4. **Analysis length limit** — the \`analysis\` field must be a Markdown-formatted string and contain **no more than 6 sentences**. Be concise. +5. **Runnable shell commands** — any index changes you recommend must be provided as **mongosh/mongo shell** commands (runnable). Use \`db.getCollection("{collectionName}")\` to reference the collection (replace \`{collectionName}\` with the actual name from \`collectionStats\`). +6. **Justify every index command** — each \`create\`/\`drop\` recommendation must include a one-sentence justification that references concrete fields/metrics from \`executionStats\` or \`indexStats\`. +7. **Prefer minimal, safe changes** — prefer a single, high-impact index over many small ones; avoid suggesting drops unless the benefit is clear and justified. +8. **Include priority** — each suggested improvement must include a \`priority\` (\`high\`/\`medium\`/\`low\`) so an engineer can triage. +9. **Be explicit about risks** — if a suggested index could increase write cost or large index size, include that as a short risk note in the improvement. +10. **Verification output** — the \`verification\` field must be a **Markdown string** (not an array). It should include one or more \`\`\`javascript code blocks\`\`\` containing **valid mongosh commands** to verify index performance or collection stats. Each command must be copy-paste runnable in mongosh (e.g. \`db.getCollection("{collectionName}").find(...).hint(...).explain("executionStats")\`). +11. **Do not change input objects** — echo input objects only under \`metadata\`; do not mutate \`{collectionStats}\`, \`{indexStats}\`, or \`{executionStats}\`—just include them as-is (and add computed helper fields if needed). +12. **If no change recommended** — return an empty \`improvements\` array and still include a short Markdown \`verification\` section to confirm the current plan. +Thinking / analysis tips (for your reasoning; do not output these tips): +- **Index-only optimization**: The best count performance occurs when all filter fields are indexed, allowing a covered query that avoids document fetches entirely. +- **Filter coverage**: Ensure all equality and range predicates in the count query are covered by an index; if not, suggest a compound index with equality fields first, range fields last. +- **COLLSCAN detection**: If totalDocsExamined is close to collection document count and no index is used, a full collection scan occurred — propose an index that minimizes this. +- **Sparse and partial indexes**: If the query filters on a field that only exists in some documents, consider a sparse or partial index to reduce index size and scan scope. +- **Equality and range ordering**: For compound indexes, equality filters should precede range filters for optimal selectivity. +- **Index-only count**: If projected or returned fields are all indexed (e.g., just counting documents matching criteria), prefer a covered plan for index-only count. +- **Write cost tradeoff**: Avoid over-indexing — recommend only indexes that materially improve count query performance or prevent full collection scans. +- If you identify indexes related to the query that have **not been accessed for a long time** or **are not selective**, consider recommending **dropping** them to reduce write and storage overhead. +- If you identify query is on a **small collection** (e.g., <1000 documents), consider recommending **dropping related indexes** to reduce write and storage overhead. +- If the **Azure_Cluster_Type** is "vCore" and a **composite index** is being created, include in \`indexOptions\` the setting: \`"storageEngine": { "enableOrderedIndex": true }\`. +Output JSON schema (required shape; adhere exactly): +\`\`\` +{ + "metadata": { + "collectionName": "", + "collectionStats": { ... }, + "indexStats": [ ... ], + "executionStats": { ... }, + "derived": { + "totalKeysExamined": , + "totalDocsExamined": , + "keysToDocsRatio": , + "usedIndex": "" + } + }, + "analysis": "", + "improvements": [ + { + "action": "create" | "drop" | "none" | "modify", + "indexSpec": { "": 1|-1, ... }, + "indexOptions": { }, + "mongoShell": "db.getCollection(\\"{collectionName}\\").createIndex({...}, {...})" , + "justification": "", + "priority": "high" | "medium" | "low", + "risks": "" + } + ], + "verification": "" +} +\`\`\` +Additional rules for the JSON: +- \`metadata.collectionName\` must be filled from \`{collectionStats.ns}\` or a suitable field; if not available set to \`null\`. +- \`derived.totalKeysExamined\`, \`derived.totalDocsExamined\`, and \`derived.keysToDocsRatio\` should be filled from \`executionStats\` if present, otherwise \`null\`. \`keysToDocsRatio\` = \`totalKeysExamined / max(1, totalDocsExamined)\`. +- \`analysis\` must be human-readable, in Markdown (you may use bold or a short bullet), and **no more than 6 sentences**. +- \`mongoShell\` commands must **only** use double quotes and valid JS object notation. +- \`verification\` must be human-readable, in Markdown. It should include one or more \`\`\`javascript code blocks\`\`\` containing valid mongosh commands. Each code block should be concise and executable as-is in mongosh. +`; + +export const CROSS_COLLECTION_QUERY_PROMPT_TEMPLATE = ` +You are an expert MongoDB assistant. Generate a MongoDB query based on the user's natural language request. +## Database Context +- **Database Name**: {databaseName} +- **User Request**: {naturalLanguageQuery} +## Available Collections and Their Schemas +{schemaInfo} + +## Query Type Requirement +- **Required Query Type**: {targetQueryType} +- You MUST generate a query of this exact type. Do not use other query types even if they might seem more appropriate. + +## Instructions +1. **Single JSON output only** — your response MUST be a single valid JSON object matching the schema below. No code fences, no surrounding text. +2. **MongoDB shell commands** — all queries must be valid MongoDB shell commands (mongosh) that can be executed directly, not javaScript functions or pseudo-code. +3. **Strict query type adherence** — you MUST generate a **{targetQueryType}** query as specified above. Ignore this requirement only if the user explicitly requests a different query type. +4. **Cross-collection queries** — the user has NOT specified a collection name, so you may need to generate queries that work across multiple collections. Consider using: + - Multiple separate queries (one per collection) if the request is collection-specific + - Aggregation pipelines with $lookup if joining data from multiple collections + - Union operations if combining results from different collections +5. **Use schema information** — examine the provided schemas to understand the data structure and field types in each collection. +6. **Respect data types** — use appropriate MongoDB operators based on the field types shown in the schema. +7. **Handle nested objects** — when you see \`type: "object"\` with \`properties\`, those are nested fields accessible with dot notation. +8. **Handle arrays** — when you see \`type: "array"\` with \`items\`, use appropriate array operators. If \`vectorLength\` is present, that's a fixed-size numeric array. +9. **Generate runnable queries** — output valid MongoDB shell syntax (mongosh) that can be executed directly. +10. **Provide clear explanation** — explain which collection(s) you're querying and why, and describe the query logic. +11. **Use db. syntax** — reference collections using \`db.collectionName\` or \`db.getCollection("collectionName")\` format. +12. **Prefer simple queries** — start with the simplest query that meets the user's needs; avoid over-complication. +13. **Consider performance** — if multiple approaches are possible, prefer the one that's more likely to be efficient. +## Query Generation Guidelines for {targetQueryType} +{queryTypeGuidelines} + +## Output JSON Schema +{outputSchema} + +## Examples +User request: "Find all users who signed up in the last 7 days" +\`\`\`json +{ + "explanation": "This query searches the 'users' collection for documents where the createdAt field is greater than or equal to 7 days ago. It uses the $gte operator to filter dates.", + "command": { + "filter": "{ \\"createdAt\\": { \\"$gte\\": { \\"$date\\": \\"<7_days_ago_ISO_string>\\" } } }", + "project": "{}", + "sort": "{}", + "skip": 0, + "limit": 0 + } +} +\`\`\` +User request: "Get total revenue by product category" +\`\`\`json +{ + "explanation": "This aggregation pipeline joins orders with products using $lookup, unwinds the product array, groups by product category, and calculates the sum of order amounts for each category, sorted by revenue descending.", + "command": { + "pipeline": "[{ \\"$lookup\\": { \\"from\\": \\"products\\", \\"localField\\": \\"productId\\", \\"foreignField\\": \\"_id\\", \\"as\\": \\"product\\" } }, { \\"$unwind\\": \\"$product\\" }, { \\"$group\\": { \\"_id\\": \\"$product.category\\", \\"totalRevenue\\": { \\"$sum\\": \\"$amount\\" } } }, { \\"$sort\\": { \\"totalRevenue\\": -1 } }]" + } +} +\`\`\` +Now generate the query based on the user's request and the provided schema information. +`; + +export const SINGLE_COLLECTION_QUERY_PROMPT_TEMPLATE = ` +You are an expert MongoDB assistant. Generate a MongoDB query based on the user's natural language request. +## Database Context +- **Database Name**: {databaseName} +- **Collection Name**: {collectionName} +- **User Request**: {naturalLanguageQuery} +## Collection Schema +{schemaInfo} +## Query Type Requirement +- **Required Query Type**: {targetQueryType} +- You MUST generate a query of this exact type. Do not use other query types even if they might seem more appropriate. + +## Instructions +1. **Single JSON output only** — your response MUST be a single valid JSON object matching the schema below. No code fences, no surrounding text. +2. **MongoDB shell commands** — all queries must be valid MongoDB shell commands (mongosh) that can be executed directly, not javaScript functions or pseudo-code. +3. **Strict query type adherence** — you MUST generate a **{targetQueryType}** query as specified above. +4. **One-sentence query** — your response must be a single, concise query that directly addresses the user's request. +5. **Return error** — When query generation is not possible (e.g., the input is invalid, contradictory, unrelated to the data schema, or incompatible with the expected query type), output an error message starts with \`Error:\` in the explanation field and \`null\` as command. +6. **Single-collection query** — the user has specified a collection name, so generate a query that works on this collection only. +7. **Use schema information** — examine the provided schema to understand the data structure and field types. +8. **Respect data types** — use appropriate MongoDB operators based on the field types shown in the schema. +9. **Handle nested objects** — when you see \`type: "object"\` with \`properties\`, those are nested fields accessible with dot notation (e.g., \`address.city\`). +10. **Handle arrays** — when you see \`type: "array"\` with \`items\`, use appropriate array operators like $elemMatch, $size, $all, etc. If \`vectorLength\` is present, that's a fixed-size numeric array (vector/embedding). +11. **Handle unions** — when you see \`type: "union"\` with \`variants\`, the field can be any of those types (handle null cases appropriately). +12. **Generate runnable queries** — output valid MongoDB shell syntax (mongosh) that can be executed directly on the specified collection. +13. **Provide clear explanation** — describe what the query does and the operators/logic used. +14. **Use db.{collectionName} syntax** — reference the collection using \`db.{collectionName}\` or \`db.getCollection("{collectionName}")\` format. +15. **Prefer simple queries** — start with the simplest query that meets the user's needs; avoid over-complication. +16. **Consider performance** — if multiple approaches are possible, prefer the one that's more likely to use indexes efficiently. +## Query Generation Guidelines for {targetQueryType} +{queryTypeGuidelines} + +## Common MongoDB Operators Reference +- **Comparison**: $eq, $ne, $gt, $gte, $lt, $lte, $in, $nin +- **Logical**: $and, $or, $not, $nor +- **Element**: $exists, $type +- **Array**: $elemMatch, $size, $all +- **Evaluation**: $regex, $text, $where, $expr +- **Aggregation**: $match, $group, $project, $sort, $limit, $lookup, $unwind +## Output JSON Schema +{outputSchema} + +## Examples +User request: "Find all documents where price is greater than 100" +\`\`\`json +{ + "explanation": "This query filters documents where the price field is greater than 100 using the $gt comparison operator.", + "command": { + "filter": "{ \\"price\\": { \\"$gt\\": 100 } }", + "project": "{}", + "sort": "{}", + "skip": 0, + "limit": 0 + } +} +\`\`\` +User request: "Get the average rating grouped by category" +\`\`\`json +{ + "explanation": "This aggregation pipeline groups documents by the category field, calculates the average rating for each group using $avg, and sorts the results by average rating in descending order.", + "command": { + "pipeline": "[{ \\"$group\\": { \\"_id\\": \\"$category\\", \\"avgRating\\": { \\"$avg\\": \\"$rating\\" } } }, { \\"$sort\\": { \\"avgRating\\": -1 } }]" + } +} +\`\`\` +User request: "Find documents with tags array containing 'featured' and status is 'active', sorted by createdAt, limit 10" +\`\`\`json +{ + "explanation": "This query finds documents where the tags array contains 'featured' and the status field equals 'active'. MongoDB's default array behavior matches any element in the array. Results are sorted by createdAt in descending order and limited to 10 documents.", + "command": { + "filter": "{ \\"tags\\": \\"featured\\", \\"status\\": \\"active\\" }", + "project": "{}", + "sort": "{ \\"createdAt\\": -1 }", + "skip": 0, + "limit": 10 + } +} +\`\`\` +Now generate the query based on the user's request and the provided collection schema. +`; + +/** + * Gets query type specific configuration (guidelines and output schema) + * @param queryType The type of query + * @returns Configuration object with guidelines and outputSchema + */ +export function getQueryTypeConfig(queryType: string): { guidelines: string; outputSchema: string } { + switch (queryType) { + case 'Find': + return { + guidelines: `- Generate a find query with appropriate filters, projections, sort, skip, and limit +- Use MongoDB query operators for filtering: $eq, $ne, $gt, $gte, $lt, $lte, $in, $nin, $regex, etc. +- Use projection to specify which fields to include or exclude +- Use sort to order results (1 for ascending, -1 for descending) +- Use skip and limit for pagination +- All fields should be valid JSON strings except skip and limit which are numbers`, + outputSchema: `\`\`\`json +{ + "explanation": "", + "command": { + "filter": "", + "project": "", + "sort": "", + "skip": , + "limit": + } +} +\`\`\``, + }; + + case 'Aggregation': + return { + guidelines: `- Generate an aggregation pipeline with appropriate stages +- Common stages: $match (filtering), $group (grouping/aggregation), $project (field selection/transformation) +- Use $sort and $limit for ordering and limiting results +- Use $lookup for joining with other collections +- Use $unwind for array expansion +- Prefer $match early in the pipeline for performance +- Pipeline should be a JSON string representing an array of stages`, + outputSchema: `\`\`\`json +{ + "explanation": "", + "command": { + "pipeline": "" + } +} +\`\`\``, + }; + + case 'Count': + return { + guidelines: `- Generate a count query with appropriate filter +- Use MongoDB query operators for filtering +- Filter should be a valid JSON string`, + outputSchema: `\`\`\`json +{ + "explanation": "", + "command": { + "filter": "" + } +} +\`\`\``, + }; + + case 'Update': + return { + guidelines: `- Generate an update query with filter and update operations +- Use update operators: $set (set field), $inc (increment), $push (add to array), $pull (remove from array) +- Specify whether to update one or many documents +- Use options like upsert if needed +- All fields should be valid JSON strings`, + outputSchema: `\`\`\`json +{ + "explanation": "", + "command": { + "filter": "", + "update": "", + "options": "" + } +} +\`\`\``, + }; + + case 'Delete': + return { + guidelines: `- Generate a delete query with appropriate filter +- Be careful with filters to avoid unintended deletions +- Filter should be a valid JSON string`, + outputSchema: `\`\`\`json +{ + "explanation": "", + "command": { + "filter": "" + } +} +\`\`\``, + }; + + default: + throw new Error(l10n.t('Unsupported query type: {queryType}', { queryType })); + } +} diff --git a/src/commands/llmEnhancedCommands/queryGenerationCommands.ts b/src/commands/llmEnhancedCommands/queryGenerationCommands.ts new file mode 100644 index 000000000..2e5b765de --- /dev/null +++ b/src/commands/llmEnhancedCommands/queryGenerationCommands.ts @@ -0,0 +1,345 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { type IActionContext } from '@microsoft/vscode-azext-utils'; +import * as l10n from '@vscode/l10n'; +import { type Document } from 'mongodb'; +import * as vscode from 'vscode'; +import { ClustersClient } from '../../documentdb/ClustersClient'; +import { ext } from '../../extensionVariables'; +import { CopilotService } from '../../services/copilotService'; +import { PromptTemplateService } from '../../services/promptTemplateService'; +import { generateSchemaDefinition, type SchemaDefinition } from '../../utils/schemaInference'; +import { FALLBACK_MODELS, PREFERRED_MODEL, getQueryTypeConfig } from './promptTemplates'; + +/** + * Type of query generation + */ +export enum QueryGenerationType { + CrossCollection = 'cross-collection', + SingleCollection = 'single-collection', +} + +/** + * Context information needed for query generation + */ +export interface QueryGenerationContext { + // The cluster/connection ID + clusterId: string; + // Database name + databaseName: string; + // Collection name (only for single-collection queries) + collectionName?: string; + // Query type of generated query + targetQueryType?: 'Find' | 'Aggregation'; + // Natural language description of the query + naturalLanguageQuery: string; + // The type of query generation + generationType: QueryGenerationType; +} + +/** + * Schema information for a collection + */ +export interface CollectionSchema { + // Collection name + collectionName: string; + // Sample documents with schema information + sampleDocuments: Array; + // Inferred schema structure + schemaDescription: string; +} + +/** + * Result from query generation + */ +export interface QueryGenerationResult { + // The generated query + generatedQuery: string; + // Explanation of the query + explanation: string; + // The model used to generate the query + modelUsed: string; +} + +/** + * Gets the prompt template for a given query generation type + * @param generationType The type of query generation + * @returns The prompt template string + */ +async function getPromptTemplate(generationType: QueryGenerationType): Promise { + return PromptTemplateService.getQueryGenerationPromptTemplate(generationType); +} + +/** + * Fills a prompt template with actual data + * @param templateType The type of template to use + * @param context The query generation context + * @param schemas Collection schemas + * @returns The filled prompt template + */ +async function fillPromptTemplate( + templateType: QueryGenerationType, + context: QueryGenerationContext, + schemas: Array, +): Promise { + // Get the template for this generation type + const template = await getPromptTemplate(templateType); + + // Determine target query type (default to Find if not specified) + const targetQueryType = context.targetQueryType || 'Find'; + + // Get query type specific guidelines and output schema + const { guidelines, outputSchema } = getQueryTypeConfig(targetQueryType); + + // Prepare schema information + let schemaInfo: string; + if (schemas.length > 0) { + if (templateType === QueryGenerationType.CrossCollection) { + schemaInfo = schemas + .map( + (schema) => + `### Collection: ${schema.collectionName || 'Unknown'}\n\nData Schema:\n\`\`\`json\n${JSON.stringify(schema.fields, null, 2)}\n\`\`\``, + ) + .join('\n\n---\n\n'); + } else { + const schema = schemas[0]; + schemaInfo = `Data Schema:\n\`\`\`json\n${JSON.stringify(schema.fields, null, 2)}\n\`\`\`\n\n`; + } + } else { + schemaInfo = `No schema information available.\n\n`; + } + + const filled = template + .replace('{databaseName}', context.databaseName) + .replace('{collectionName}', context.collectionName || 'N/A') + .replace(/{targetQueryType}/g, targetQueryType) + .replace('{queryTypeGuidelines}', guidelines) + .replace('{outputSchema}', outputSchema) + .replace('{schemaInfo}', schemaInfo) + .replace('{naturalLanguageQuery}', context.naturalLanguageQuery); + + return filled; +} + +/** + * Generates a MongoDB query based on natural language input + * @param context Action context for telemetry + * @param queryContext Query generation context + * @returns Generated query and explanation + */ +export async function generateQuery( + context: IActionContext, + queryContext: QueryGenerationContext, +): Promise { + ext.outputChannel.trace( + l10n.t('Query generation started: type={type}, targetQueryType={targetQueryType}', { + type: queryContext.generationType, + targetQueryType: queryContext.targetQueryType || 'Find', + }), + ); + + // Check if Copilot is available + ext.outputChannel.trace(l10n.t('Checking GitHub Copilot availability...')); + const copilotAvailable = await CopilotService.isAvailable(); + if (!copilotAvailable) { + ext.outputChannel.error(l10n.t('GitHub Copilot is not available')); + throw new Error( + l10n.t( + 'GitHub Copilot is not available. Please install the GitHub Copilot extension and ensure you have an active subscription.', + ), + ); + } + ext.outputChannel.trace(l10n.t('GitHub Copilot is available')); + + // Get the MongoDB client + ext.outputChannel.trace(l10n.t('Connecting to cluster: {clusterId}', { clusterId: queryContext.clusterId })); + const getClientStart = Date.now(); + const client = await ClustersClient.getClient(queryContext.clusterId); + context.telemetry.measurements.getClientDurationMs = Date.now() - getClientStart; + ext.outputChannel.trace( + l10n.t('Connected to cluster [{durationMs}ms]', { + durationMs: context.telemetry.measurements.getClientDurationMs.toString(), + }), + ); + + // Gather schema information + ext.outputChannel.trace(l10n.t('Gathering schema information...')); + const schemas: Array = []; + const schemaGatheringStart = Date.now(); + + try { + if (queryContext.generationType === QueryGenerationType.CrossCollection) { + ext.outputChannel.trace( + l10n.t('Discovering collections in database: {databaseName}', { + databaseName: queryContext.databaseName, + }), + ); + // Get all collections in the database + const listCollectionsStart = Date.now(); + const collections = await client.listCollections(queryContext.databaseName); + context.telemetry.measurements.listCollectionsDurationMs = Date.now() - listCollectionsStart; + ext.outputChannel.trace( + l10n.t('Found {count} collections [{durationMs}ms]', { + count: collections.length.toString(), + durationMs: context.telemetry.measurements.listCollectionsDurationMs.toString(), + }), + ); + + let collectionIndex = 0; + for (const collection of collections) { + collectionIndex++; + ext.outputChannel.trace( + l10n.t('Sampling documents from collection: {collectionName}', { + collectionName: collection.name, + }), + ); + const sampleDocsStart = Date.now(); + const sampleDocs = await client.getSampleDocuments(queryContext.databaseName, collection.name, 3); + const sampleDocsDuration = Date.now() - sampleDocsStart; + context.telemetry.measurements[`sampleDocs_${collectionIndex}_DurationMs`] = sampleDocsDuration; + ext.outputChannel.trace( + l10n.t('Sampled {count} documents from collection: {collectionName} [{durationMs}ms]', { + count: sampleDocs.length.toString(), + collectionName: collection.name, + durationMs: sampleDocsDuration.toString(), + }), + ); + + if (sampleDocs.length > 0) { + ext.outputChannel.trace( + l10n.t('Generating schema from {count} sample documents for collection: {collectionName}', { + count: sampleDocs.length.toString(), + collectionName: collection.name, + }), + ); + const schema = generateSchemaDefinition(sampleDocs, collection.name); + schemas.push(schema); + } else { + ext.outputChannel.trace( + l10n.t('No sample documents found for collection: {collectionName}', { + collectionName: collection.name, + }), + ); + schemas.push({ collectionName: collection.name, fields: {} }); + } + } + } else { + if (!queryContext.collectionName) { + throw new Error(l10n.t('Collection name is required for single-collection query generation')); + } + + ext.outputChannel.trace( + l10n.t('Sampling documents from collection: {databaseName}.{collectionName}', { + databaseName: queryContext.databaseName, + collectionName: queryContext.collectionName, + }), + ); + const sampleDocsStart = Date.now(); + const sampleDocs = await client.getSampleDocuments( + queryContext.databaseName, + queryContext.collectionName, + 10, + ); + context.telemetry.measurements.sampleDocsDurationMs = Date.now() - sampleDocsStart; + ext.outputChannel.trace( + l10n.t('Retrieved {count} sample documents [{durationMs}ms]', { + count: sampleDocs.length.toString(), + durationMs: context.telemetry.measurements.sampleDocsDurationMs.toString(), + }), + ); + + ext.outputChannel.trace(l10n.t('Generating schema from sample documents...')); + const schema = generateSchemaDefinition(sampleDocs, queryContext.collectionName); + schemas.push(schema); + } + ext.outputChannel.trace( + l10n.t('Schema generation complete for {count} collection(s)', { count: schemas.length.toString() }), + ); + context.telemetry.measurements.schemaGatheringDurationMs = Date.now() - schemaGatheringStart; + ext.outputChannel.trace( + l10n.t('Schema gathering completed [{durationMs}ms]', { + durationMs: context.telemetry.measurements.schemaGatheringDurationMs.toString(), + }), + ); + } catch (error) { + context.telemetry.measurements.schemaGatheringDurationMs = Date.now() - schemaGatheringStart; + ext.outputChannel.error( + l10n.t('Schema gathering failed: {message}', { + message: error instanceof Error ? error.message : String(error), + }), + ); + throw new Error( + l10n.t('Failed to gather schema information: {message}', { + message: error instanceof Error ? error.message : String(error), + }), + ); + } + + // Fill the prompt template + ext.outputChannel.trace(l10n.t('Filling prompt template...')); + const promptContent = await fillPromptTemplate(queryContext.generationType, queryContext, schemas); + ext.outputChannel.trace(l10n.t('Prompt template filled successfully')); + + // Send to Copilot with configured models + ext.outputChannel.trace( + l10n.t('Sending request to GitHub Copilot (preferred model: {preferredModel})...', { + preferredModel: PREFERRED_MODEL || 'none specified', + }), + ); + const llmCallStart = Date.now(); + const response = await CopilotService.sendMessage([vscode.LanguageModelChatMessage.User(promptContent)], { + preferredModel: PREFERRED_MODEL, + fallbackModels: FALLBACK_MODELS, + }); + context.telemetry.measurements.llmCallDurationMs = Date.now() - llmCallStart; + ext.outputChannel.trace( + l10n.t('Received response from model: {modelUsed} [{durationMs}ms]', { + modelUsed: response.modelUsed, + durationMs: context.telemetry.measurements.llmCallDurationMs.toString(), + }), + ); + + // Check if the preferred model was used + if (response.modelUsed !== PREFERRED_MODEL && PREFERRED_MODEL) { + // Show warning if not using preferred model + void vscode.window.showWarningMessage( + l10n.t( + 'Query generation is using model "{actualModel}" instead of preferred "{preferredModel}". Results may vary.', + { + actualModel: response.modelUsed, + preferredModel: PREFERRED_MODEL, + }, + ), + ); + } + + // Add telemetry for the model used + context.telemetry.properties.modelUsed = response.modelUsed; + context.telemetry.properties.generationType = queryContext.targetQueryType || 'Find'; + + // Parse the response + ext.outputChannel.trace(l10n.t('Parsing language model response...')); + try { + const result = JSON.parse(response.text) as { explanation: string; command: Record }; + if (result.command === undefined || result.command === null || result.explanation.startsWith('Error:')) { + throw new Error(result.explanation); + } + + ext.outputChannel.trace(l10n.t('Query generation completed successfully')); + return { + generatedQuery: JSON.stringify(result.command, null, 2), + explanation: result.explanation, + modelUsed: response.modelUsed, + }; + } catch { + ext.outputChannel.error(l10n.t('Failed to parse language model response')); + throw new Error( + l10n.t('Failed to parse the response from the language model. LLM output:\n{output}', { + output: response.text, + }), + ); + } +} diff --git a/src/documentdb/ClustersClient.ts b/src/documentdb/ClustersClient.ts index 0293040ad..8f2759ab9 100644 --- a/src/documentdb/ClustersClient.ts +++ b/src/documentdb/ClustersClient.ts @@ -35,6 +35,16 @@ import { AuthMethodId } from './auth/AuthMethod'; import { MicrosoftEntraIDAuthHandler } from './auth/MicrosoftEntraIDAuthHandler'; import { NativeAuthHandler } from './auth/NativeAuthHandler'; import { CredentialCache, type CachedClusterCredentials } from './CredentialCache'; +import { + llmEnhancedFeatureApis, + type CollectionStats, + type CreateIndexResult, + type DropIndexResult, + type ExplainOptions, + type ExplainResult, + type IndexSpecification, + type IndexStats, +} from './LlmEnhancedFeatureApis'; import { getHostsFromConnectionString, hasAzureDomain } from './utils/connectionStringHelpers'; import { getClusterMetadata, type ClusterMetadata } from './utils/getClusterMetadata'; import { toFilterQueryObj } from './utils/toFilterQuery'; @@ -106,6 +116,8 @@ export class ClustersClient { static _clients: Map = new Map(); private _mongoClient: MongoClient; + private _llmEnhancedFeatureApis: llmEnhancedFeatureApis | null = null; + private _clusterMetadataPromise: Promise | null = null; /** * Use getClient instead of a constructor. Connections/Client are being cached and reused. @@ -175,10 +187,12 @@ export class ClustersClient { // Connect with the configured options await this.connect(connectionString, options, credentials.emulatorConfiguration); - // Collect telemetry (non-blocking) - void callWithTelemetryAndErrorHandling('connect.getmetadata', async (context) => { - const metadata: ClusterMetadata = await getClusterMetadata(this._mongoClient, hosts); + // Start metadata collection and store the promise + this._clusterMetadataPromise = getClusterMetadata(this._mongoClient, hosts); + // Collect telemetry (non-blocking) - reuses the same promise + void callWithTelemetryAndErrorHandling('connect.getmetadata', async (context) => { + const metadata: ClusterMetadata = await this._clusterMetadataPromise!; context.telemetry.properties = { authmethod: authMethod, ...context.telemetry.properties, @@ -194,6 +208,7 @@ export class ClustersClient { ): Promise { try { this._mongoClient = await MongoClient.connect(connectionString, options); + this._llmEnhancedFeatureApis = new llmEnhancedFeatureApis(this._mongoClient); } catch (error) { const message = parseError(error).message; if (emulatorConfiguration?.isEmulator && message.includes('ECONNREFUSED')) { @@ -230,6 +245,7 @@ export class ClustersClient { await client._mongoClient.connect(); } else { client = new ClustersClient(credentialId); + // Cluster metadata is set in initClient await client.initClient(); ClustersClient._clients.set(credentialId, client); } @@ -237,6 +253,21 @@ export class ClustersClient { return client; } + /** + * Retrieves cluster metadata for this client instance. + * + * @returns A promise that resolves to cluster metadata. + */ + public async getClusterMetadata(): Promise { + if (this._clusterMetadataPromise) { + return this._clusterMetadataPromise; + } + + // This should not happen as the promise is initialized in initClient, + // but if it does, we throw an error rather than trying to recover + throw new Error(l10n.t('Cluster metadata not initialized. Client may not be properly connected.')); + } + /** * Determines whether a client for the given credential identifier is present in the internal cache. */ @@ -642,4 +673,128 @@ export class ClustersClient { }; } } + + // ========================================== + // LLM Enhanced Feature APIs + // ========================================== + + /** + * Get detailed index statistics for a collection + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @returns Array of index statistics including usage information + */ + async getIndexStats(databaseName: string, collectionName: string): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.getIndexStats(databaseName, collectionName); + } + + /** + * Get detailed collection statistics + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @returns Collection statistics including size, count, and index information + */ + async getCollectionStats(databaseName: string, collectionName: string): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.getCollectionStats(databaseName, collectionName); + } + + /** + * Explain a find query with full execution statistics + * Supports sort, projection, skip, and limit options + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param options - Query options including filter, sort, projection, skip, and limit + * @returns Detailed explain result with execution statistics + */ + async explainFind( + databaseName: string, + collectionName: string, + options: ExplainOptions = {}, + ): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.explainFind(databaseName, collectionName, options); + } + + /** + * Explain an aggregation pipeline with full execution statistics + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param pipeline - Aggregation pipeline stages + * @returns Detailed explain result with execution statistics + */ + async explainAggregate(databaseName: string, collectionName: string, pipeline: Document[]): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.explainAggregate(databaseName, collectionName, pipeline); + } + + /** + * Explain a count operation with full execution statistics + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param filter - Query filter for the count operation + * @returns Detailed explain result with execution statistics + */ + async explainCount(databaseName: string, collectionName: string, filter: Filter = {}): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.explainCount(databaseName, collectionName, filter); + } + + /** + * Create an index on a collection + * Supports both simple and composite indexes with various options + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param indexSpec - Index specification including key and options + * @returns Result of the index creation operation + */ + async createIndex( + databaseName: string, + collectionName: string, + indexSpec: IndexSpecification, + ): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.createIndex(databaseName, collectionName, indexSpec); + } + + /** + * Drop an index from a collection + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param indexName - Name of the index to drop (use "*" to drop all non-_id indexes) + * @returns Result of the index drop operation + */ + async dropIndex(databaseName: string, collectionName: string, indexName: string): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.dropIndex(databaseName, collectionName, indexName); + } + + /** + * Get sample documents from a collection using random sampling + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param limit - Maximum number of documents to sample (default: 10) + * @returns Array of sample documents + */ + async getSampleDocuments(databaseName: string, collectionName: string, limit: number = 10): Promise { + if (!this._llmEnhancedFeatureApis) { + throw new Error('LLM Enhanced Feature APIs not initialized. Ensure the client is connected.'); + } + return this._llmEnhancedFeatureApis.getSampleDocuments(databaseName, collectionName, limit); + } } diff --git a/src/documentdb/LlmEnhancedFeatureApis.ts b/src/documentdb/LlmEnhancedFeatureApis.ts new file mode 100644 index 000000000..aaee87cb1 --- /dev/null +++ b/src/documentdb/LlmEnhancedFeatureApis.ts @@ -0,0 +1,429 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * LLM Enhanced Feature APIs + */ + +import { type Document, type Filter, type MongoClient, type Sort } from 'mongodb'; + +/** + * Options for explain operations + */ +export interface ExplainOptions { + // The query filter + filter?: Filter; + // Sort specification + sort?: Sort; + // Projection specification + projection?: Document; + // Number of documents to skip + skip?: number; + // Maximum number of documents to return + limit?: number; +} + +/** + * Index specification for creating indexes + * Supports both simple and composite indexes + */ +export interface IndexSpecification { + // Index key specification + key: Record; + // Index name + name?: string; + // Create a unique index + unique?: boolean; + // Create index in the background + background?: boolean; + // Create a sparse index + sparse?: boolean; + // TTL for documents in seconds + expireAfterSeconds?: number; + // Partial index filter expression + partialFilterExpression?: Document; + // Additional index options + [key: string]: unknown; +} + +/** + * Result of index creation operation + */ +export interface CreateIndexResult { + // Operation status + ok: number; + // Name of the created index + indexName?: string; + // Number of indexes after creation + numIndexesAfter?: number; + // Number of indexes before creation + numIndexesBefore?: number; + // Notes or warnings + note?: string; +} + +/** + * Result of index drop operation + */ +export interface DropIndexResult { + // Operation status + ok: number; + // Number of indexes after dropping + nIndexesWas?: number; +} + +/** + * Collection statistics result + */ +export interface CollectionStats { + // Namespace (database.collection) + ns: string; + // Number of documents in the collection + count: number; + // Total size of all documents in bytes + size: number; + // Average object size in bytes + avgObjSize: number; + // Storage size in bytes + storageSize: number; + // Number of indexes + nindexes: number; + // Total index size in bytes + totalIndexSize: number; + // Individual index sizes + indexSizes: Record; +} + +/** + * Index statistics for a single index + */ +export interface IndexStats { + // Index name + name: string; + // Index key specification + key: Record; + // Host information + host: string; + // Access statistics + accesses: { + // Number of times the index has been used + ops: number; + // Timestamp of last access + since: Date; + }; +} + +/** + * Explain plan result with execution statistics + */ +export interface ExplainResult { + // Query planner information + queryPlanner: { + // MongoDB version + mongodbVersion?: string; + // Namespace + namespace: string; + // Whether index was used + indexFilterSet: boolean; + // Parsed query + parsedQuery?: Document; + // Winning plan + winningPlan: Document; + // Rejected plans + rejectedPlans?: Document[]; + }; + // Execution statistics + executionStats?: { + // Execution success status + executionSuccess: boolean; + // Number of documents returned + nReturned: number; + // Execution time in milliseconds + executionTimeMillis: number; + // Total number of keys examined + totalKeysExamined: number; + // Total number of documents examined + totalDocsExamined: number; + // Detailed execution stages + executionStages: Document; + }; + // Server information + serverInfo?: { + host: string; + port: number; + version: string; + }; + // Operation status + ok: number; +} + +/** + * LLM Enhanced Feature APIs + */ +export class llmEnhancedFeatureApis { + constructor(private readonly mongoClient: MongoClient) {} + + /** + * Get statistics for all indexes in a collection + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @returns Array of index statistics + */ + async getIndexStats(databaseName: string, collectionName: string): Promise { + const collection = this.mongoClient.db(databaseName).collection(collectionName); + + const indexStatsResult = await collection + .aggregate([ + { + $indexStats: {}, + }, + ]) + .toArray(); + + return indexStatsResult.map((stat) => { + const accesses = stat.accesses as { ops?: number; since?: Date } | undefined; + + return { + name: stat.name as string, + key: stat.key as Record, + host: stat.host as string, + accesses: { + ops: accesses?.ops ?? 0, + since: accesses?.since ?? new Date(), + }, + }; + }); + } + + /** + * Get statistics for a collection + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @returns Collection statistics + */ + async getCollectionStats(databaseName: string, collectionName: string): Promise { + const db = this.mongoClient.db(databaseName); + + // Use the collStats command to get detailed collection statistics + const stats = await db.command({ + collStats: collectionName, + }); + + return { + ns: stats.ns as string, + count: (stats.count as number) ?? 0, + size: (stats.size as number) ?? 0, + avgObjSize: (stats.avgObjSize as number) ?? 0, + storageSize: (stats.storageSize as number) ?? 0, + nindexes: (stats.nindexes as number) ?? 0, + totalIndexSize: (stats.totalIndexSize as number) ?? 0, + indexSizes: (stats.indexSizes as Record) ?? {}, + }; + } + + /** + * Explain a find query with full execution statistics + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param options - Query options including filter, sort, projection, skip, and limit + * @returns Detailed explain result with execution statistics + */ + async explainFind( + databaseName: string, + collectionName: string, + options: ExplainOptions = {}, + ): Promise { + const db = this.mongoClient.db(databaseName); + + const { filter = {}, sort, projection, skip, limit } = options; + + const findCmd: Document = { + find: collectionName, + filter, + }; + + // Add optional fields if they are defined + if (sort !== undefined) { + findCmd.sort = sort; + } + + if (projection !== undefined) { + findCmd.projection = projection; + } + + if (skip !== undefined && skip >= 0) { + findCmd.skip = skip; + } + + if (limit !== undefined && limit >= 0) { + findCmd.limit = limit; + } + + const command: Document = { + explain: findCmd, + verbosity: 'executionStats', + }; + + const explainResult = await db.command(command); + + return explainResult as ExplainResult; + } + + /** + * Explain an aggregation pipeline with full execution statistics + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param pipeline - Aggregation pipeline stages + * @returns Detailed explain result with execution statistics + */ + async explainAggregate(databaseName: string, collectionName: string, pipeline: Document[]): Promise { + const db = this.mongoClient.db(databaseName); + + const command: Document = { + explain: { + aggregate: collectionName, + pipeline, + cursor: {}, + }, + verbosity: 'executionStats', + }; + + const explainResult = await db.command(command); + + return explainResult as ExplainResult; + } + + /** + * Explain a count operation with full execution statistics + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param filter - Query filter for the count operation + * @returns Detailed explain result with execution statistics + */ + async explainCount(databaseName: string, collectionName: string, filter: Filter = {}): Promise { + const db = this.mongoClient.db(databaseName); + + const command: Document = { + explain: { + count: collectionName, + query: filter, + }, + verbosity: 'executionStats', + }; + + const explainResult = await db.command(command); + + return explainResult; + } + + /** + * Create an index on a collection + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param indexSpec - Index specification including key and options + * @returns Result of the index creation operation + */ + async createIndex( + databaseName: string, + collectionName: string, + indexSpec: IndexSpecification, + ): Promise { + const db = this.mongoClient.db(databaseName); + + const { key, name, unique, background, sparse, expireAfterSeconds, partialFilterExpression, ...otherOptions } = + indexSpec; + + const indexDefinition: Document = { + key, + }; + + // Add optional fields only if they are defined + if (name !== undefined) { + indexDefinition.name = name; + } + + if (unique !== undefined) { + indexDefinition.unique = unique; + } + + if (background !== undefined) { + indexDefinition.background = background; + } + + if (sparse !== undefined) { + indexDefinition.sparse = sparse; + } + + if (expireAfterSeconds !== undefined) { + indexDefinition.expireAfterSeconds = expireAfterSeconds; + } + + if (partialFilterExpression !== undefined) { + indexDefinition.partialFilterExpression = partialFilterExpression; + } + + // Add any other options + Object.assign(indexDefinition, otherOptions); + + const command: Document = { + createIndexes: collectionName, + indexes: [indexDefinition], + }; + + const result = await db.command(command); + + return { + ok: (result.ok as number) ?? 0, + indexName: result.createdCollectionAutomatically !== undefined ? (name ?? 'auto-generated') : undefined, + numIndexesAfter: result.numIndexesAfter as number | undefined, + numIndexesBefore: result.numIndexesBefore as number | undefined, + note: result.note as string | undefined, + }; + } + + /** + * Drop an index from a collection + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param indexName - Name of the index to drop (use "*" to drop all non-_id indexes) + * @returns Result of the index drop operation + */ + async dropIndex(databaseName: string, collectionName: string, indexName: string): Promise { + const db = this.mongoClient.db(databaseName); + + const command: Document = { + dropIndexes: collectionName, + index: indexName, + }; + + const result = await db.command(command); + + return { + ok: (result.ok as number) ?? 0, + nIndexesWas: result.nIndexesWas as number | undefined, + }; + } + + /** + * Get sample documents from a collection using random sampling + * @param databaseName - Name of the database + * @param collectionName - Name of the collection + * @param limit - Maximum number of documents to sample (default: 10) + * @returns Array of sample documents + */ + async getSampleDocuments(databaseName: string, collectionName: string, limit: number = 10): Promise { + const collection = this.mongoClient.db(databaseName).collection(collectionName); + + const sampleDocuments = await collection + .aggregate([ + { + $sample: { size: limit }, + }, + ]) + .toArray(); + + return sampleDocuments; + } +} diff --git a/src/services/copilotService.ts b/src/services/copilotService.ts new file mode 100644 index 000000000..11cac68f7 --- /dev/null +++ b/src/services/copilotService.ts @@ -0,0 +1,151 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as l10n from '@vscode/l10n'; +import * as vscode from 'vscode'; + +/** + * Options for sending a message to the language model + */ +export interface CopilotMessageOptions { + /* The preferred model to use + If the specified model is not available, will fall back to available models + */ + preferredModel?: string; + + /* List of fallback models */ + fallbackModels?: string[]; + + // TODO: + /* Temperature setting for the model (if supported later) */ + // temperature?: number; + + /* Maximum tokens for the response (if supported later) */ + // maxTokens?: number; +} + +/** + * Response from the Copilot service + */ +export interface CopilotResponse { + /* The generated text response */ + text: string; + /* The model used to generate the response */ + modelUsed: string; +} + +/** + * Service for interacting with GitHub Copilot's LLM + */ +export class CopilotService { + /** + * Sends a message to the Copilot LLM and returns the response + * + * @param messages - Array of chat messages to send to the model + * @param options - Options for the request + * @returns The response from the model + * @throws Error if no suitable model is available or if the user cancels + */ + static async sendMessage( + messages: vscode.LanguageModelChatMessage[], + options?: CopilotMessageOptions, + ): Promise { + // Get all available models from VS Code + const availableModels = await vscode.lm.selectChatModels({ vendor: 'copilot' }); + + const preferredModels = this.getPreferredModels(options); + const selectedModel = this.selectBestModel(availableModels, preferredModels); + + if (!selectedModel) { + throw new Error( + l10n.t( + 'No suitable language model found. Please ensure GitHub Copilot is installed and you have an active subscription.', + ), + ); + } + + const response = await this.sendToModel(selectedModel, messages, options); + return { + text: response, + modelUsed: selectedModel.id, + }; + } + + /** + * Builds the ordered list of preferred models + */ + private static getPreferredModels(options?: CopilotMessageOptions): string[] { + const models: string[] = []; + + if (options?.preferredModel) { + models.push(options.preferredModel); + } + + if (options?.fallbackModels && options.fallbackModels.length > 0) { + models.push(...options.fallbackModels); + } + + return models; + } + + /** + * Selects the best available model based on preference order + * + * @param availableModels - All available models from VS Code + * @param preferredModels - Ordered list of preferred model families + * @returns The best matching model, or the first available model if no matches + */ + private static selectBestModel( + availableModels: vscode.LanguageModelChat[], + preferredModels: string[], + ): vscode.LanguageModelChat | undefined { + if (availableModels.length === 0) { + return undefined; + } + + if (preferredModels.length !== 0) { + for (const preferredModel of preferredModels) { + const matchingModel = availableModels.find((model) => model.id === preferredModel); + if (matchingModel) { + return matchingModel; + } + } + } + return availableModels[0]; + } + + /** + * Sends messages to a specific model and collects the response + */ + private static async sendToModel( + model: vscode.LanguageModelChat, + messages: vscode.LanguageModelChatMessage[], + _options?: CopilotMessageOptions, + ): Promise { + // Github copilot LLM API currently doesn't support temperature or maxTokens in + // LanguageModelChatRequestOptions, but we keep them here for potential future use + const requestOptions: vscode.LanguageModelChatRequestOptions = {}; + + const chatResponse = await model.sendRequest(messages, requestOptions); + + // Collect the streaming response + let fullResponse = ''; + for await (const fragment of chatResponse.text) { + fullResponse += fragment; + } + + return fullResponse; + } + + /** + * Checks if LLMs are available + * + * @returns true if at least one model is available + */ + static async isAvailable(): Promise { + const models = await vscode.lm.selectChatModels({ vendor: 'copilot' }); + return models.length > 0; + } +} diff --git a/src/services/promptTemplateService.ts b/src/services/promptTemplateService.ts new file mode 100644 index 000000000..cca23965f --- /dev/null +++ b/src/services/promptTemplateService.ts @@ -0,0 +1,245 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as l10n from '@vscode/l10n'; +import * as fs from 'fs/promises'; +import * as path from 'path'; +import * as vscode from 'vscode'; +import { + CROSS_COLLECTION_QUERY_PROMPT_TEMPLATE, + SINGLE_COLLECTION_QUERY_PROMPT_TEMPLATE, +} from '../commands/llmEnhancedCommands/promptTemplates'; +import { QueryGenerationType } from '../commands/llmEnhancedCommands/queryGenerationCommands'; + +/** + * Service for loading prompt templates from custom files or built-in templates + */ +export class PromptTemplateService { + private static readonly configSection = 'documentDB.aiAssistant'; + // private static readonly templateCache: Map = new Map(); + private static readonly templateCache: Map = new Map(); + + // /** + // * Gets the prompt template for index advisor + // * @param commandType The type of command (find, aggregate, or count) + // * @returns The prompt template string + // */ + // public static async getIndexAdvisorPromptTemplate(commandType: CommandType): Promise { + // // Get configuration + // const config = vscode.workspace.getConfiguration(this.configSection); + // const cacheEnabled = config.get('enablePromptCache', true); + + // // Check if have a cached template + // if (cacheEnabled) { + // const cached = this.templateCache.get(commandType); + // if (cached) { + // return cached; + // } + // } + + // // Get the configuration key for this command type + // const configKey = this.getIndexAdvisorConfigKey(commandType); + + // // Check if a custom template path is configured + // const customTemplatePath = config.get(configKey); + + // let template: string; + + // if (customTemplatePath) { + // try { + // // Load custom template from file + // template = await this.loadTemplateFromFile(customTemplatePath, commandType.toString()); + // void vscode.window.showInformationMessage( + // l10n.t('Using custom prompt template for {type} query: {path}', { + // type: commandType, + // path: customTemplatePath, + // }), + // ); + // } catch (error) { + // // Log error and fall back to built-in template + // void vscode.window.showWarningMessage( + // l10n.t('Failed to load custom prompt template from {path}: {error}. Using built-in template.', { + // path: customTemplatePath, + // error: error instanceof Error ? error.message : String(error), + // }), + // ); + // template = this.getBuiltInIndexAdvisorTemplate(commandType); + // } + // } else { + // // Use built-in template + // template = this.getBuiltInIndexAdvisorTemplate(commandType); + // } + + // // Cache the template (if caching is enabled) + // if (cacheEnabled) { + // this.templateCache.set(commandType, template); + // } + + // return template; + // } + + /** + * Gets the prompt template for query generation + * @param generationType The type of query generation (cross-collection or single-collection) + * @returns The prompt template string + */ + public static async getQueryGenerationPromptTemplate(generationType: QueryGenerationType): Promise { + // Get configuration + const config = vscode.workspace.getConfiguration(this.configSection); + const configKey = this.getQueryGenerationConfigKey(generationType); + const customTemplatePath = config.get(configKey); + + let template: string; + + if (customTemplatePath) { + try { + template = await this.loadTemplateFromFile(customTemplatePath, generationType.toString()); + void vscode.window.showInformationMessage( + l10n.t('Using custom prompt template for {type} query generation: {path}', { + type: generationType, + path: customTemplatePath, + }), + ); + } catch (error) { + void vscode.window.showWarningMessage( + l10n.t('Failed to load custom prompt template from {path}: {error}. Using built-in template.', { + path: customTemplatePath, + error: error instanceof Error ? error.message : String(error), + }), + ); + template = this.getBuiltInQueryGenerationTemplate(generationType); + } + } else { + // Use built-in template + template = this.getBuiltInQueryGenerationTemplate(generationType); + } + + return template; + } + + /** + * Clears the template cache, forcing templates to be reloaded on next use + */ + public static clearCache(): void { + this.templateCache.clear(); + } + + // /** + // * Gets the configuration key for a command type + // * @param commandType The command type + // * @returns The configuration key + // */ + // private static getIndexAdvisorConfigKey(commandType: CommandType): string { + // switch (commandType) { + // case CommandType.Find: + // return 'findQueryPromptPath'; + // case CommandType.Aggregate: + // return 'aggregateQueryPromptPath'; + // case CommandType.Count: + // return 'countQueryPromptPath'; + // default: + // throw new Error(l10n.t('Unknown command type: {type}', { type: commandType })); + // } + // } + + /** + * Gets the configuration key for a query generation type + * @param generationType The query generation type + * @returns The configuration key + */ + private static getQueryGenerationConfigKey(generationType: QueryGenerationType): string { + switch (generationType) { + case QueryGenerationType.CrossCollection: + return 'crossCollectionQueryPromptPath'; + case QueryGenerationType.SingleCollection: + return 'singleCollectionQueryPromptPath'; + default: + throw new Error(l10n.t('Unknown query generation type: {type}', { type: generationType })); + } + } + + // /** + // * Gets the built-in prompt template for a command type + // * @param commandType The command type + // * @returns The built-in template + // */ + // private static getBuiltInIndexAdvisorTemplate(commandType: CommandType): string { + // switch (commandType) { + // case CommandType.Find: + // return FIND_QUERY_PROMPT_TEMPLATE; + // case CommandType.Aggregate: + // return AGGREGATE_QUERY_PROMPT_TEMPLATE; + // case CommandType.Count: + // return COUNT_QUERY_PROMPT_TEMPLATE; + // default: + // throw new Error(l10n.t('Unknown command type: {type}', { type: commandType })); + // } + // } + + /** + * Gets the built-in prompt template for a query generation type + * @param generationType The query generation type + * @returns The built-in template + */ + private static getBuiltInQueryGenerationTemplate(generationType: QueryGenerationType): string { + switch (generationType) { + case QueryGenerationType.CrossCollection: + return CROSS_COLLECTION_QUERY_PROMPT_TEMPLATE; + case QueryGenerationType.SingleCollection: + return SINGLE_COLLECTION_QUERY_PROMPT_TEMPLATE; + default: + throw new Error(l10n.t('Unknown query generation type: {type}', { type: generationType })); + } + } + + /** + * Loads a template from a file path + * @param filePath The absolute or relative file path + * @param templateType The template type identifier (for error messages) + * @returns The template content + */ + private static async loadTemplateFromFile(filePath: string, templateType: string): Promise { + try { + let resolvedPath = filePath; + if (!path.isAbsolute(filePath)) { + const workspaceFolders = vscode.workspace.workspaceFolders; + if (workspaceFolders && workspaceFolders.length > 0) { + resolvedPath = path.join(workspaceFolders[0].uri.fsPath, filePath); + } + } + + // Check if file exists + try { + await fs.access(resolvedPath); + } catch { + throw new Error( + l10n.t('Template file not found: {path}', { + path: resolvedPath, + }), + ); + } + + // Read the file + const content = await fs.readFile(resolvedPath, 'utf-8'); + + if (!content || content.trim().length === 0) { + throw new Error( + l10n.t('Template file is empty: {path}', { + path: resolvedPath, + }), + ); + } + + return content; + } catch (error) { + throw new Error( + l10n.t('Failed to load template file for {type}: {error}', { + type: templateType, + error: error instanceof Error ? error.message : String(error), + }), + ); + } + } +} diff --git a/src/utils/schemaInference.ts b/src/utils/schemaInference.ts new file mode 100644 index 000000000..b7226affd --- /dev/null +++ b/src/utils/schemaInference.ts @@ -0,0 +1,339 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Binary, Decimal128, Double, Int32, Long, ObjectId, Timestamp, type Document } from 'mongodb'; + +type PrimitiveType = + | 'string' + | 'number' + | 'boolean' + | 'objectId' + | 'date' + | 'binary' + | 'regex' + | 'timestamp' + | 'undefined' + | 'unknown'; + +interface FieldSummary { + primitiveTypes: Set; + hasNull: boolean; + objectProperties?: Map; + arraySummary?: ArraySummary; +} + +interface ArraySummary { + elementSummary: FieldSummary; + vectorLengths: Set; + nonVectorObserved: boolean; + sawValues: boolean; +} + +export interface SchemaDefinition { + collectionName?: string; + fields: Record; +} + +export type SchemaFieldDefinition = string | SchemaObjectDefinition | SchemaArrayDefinition | SchemaUnionDefinition; + +export interface SchemaObjectDefinition { + type: 'object'; + properties: Record; +} + +export interface SchemaArrayDefinition { + type: 'array'; + items: SchemaFieldDefinition; + vectorLength?: number; +} + +export interface SchemaUnionDefinition { + type: 'union'; + variants: SchemaFieldDefinition[]; +} + +export function generateSchemaDefinition(documents: Array, collectionName?: string): SchemaDefinition { + const root = new Map(); + + for (const doc of documents) { + recordDocument(root, doc); + } + + const fields = convertProperties(root); + + const schema: SchemaDefinition = { fields }; + if (collectionName) { + schema.collectionName = collectionName; + } + + return schema; +} + +function recordDocument(target: Map, doc: Document): void { + for (const [key, value] of Object.entries(doc)) { + const summary = target.get(key) ?? createFieldSummary(); + recordValue(summary, value); + target.set(key, summary); + } +} + +function recordValue(summary: FieldSummary, value: unknown): void { + if (value === null) { + summary.hasNull = true; + return; + } + + if (Array.isArray(value)) { + handleArray(summary, value); + return; + } + + if (isPlainObject(value)) { + handleObject(summary, value as Record); + return; + } + + summary.primitiveTypes.add(getPrimitiveType(value)); +} + +function handleObject(summary: FieldSummary, value: Record): void { + const properties = summary.objectProperties ?? new Map(); + summary.objectProperties = properties; + + for (const [key, nested] of Object.entries(value)) { + const nestedSummary = properties.get(key) ?? createFieldSummary(); + recordValue(nestedSummary, nested); + properties.set(key, nestedSummary); + } +} + +function handleArray(summary: FieldSummary, values: Array): void { + const arraySummary = summary.arraySummary ?? createArraySummary(); + summary.arraySummary = arraySummary; + + arraySummary.sawValues ||= values.length > 0; + const vectorCandidate = values.length > 0 && values.every((element) => isNumericValue(element)); + + if (vectorCandidate) { + arraySummary.vectorLengths.add(values.length); + } else if (values.length > 0) { + arraySummary.nonVectorObserved = true; + } + + for (const element of values) { + recordValue(arraySummary.elementSummary, element); + } +} + +function getPrimitiveType(value: unknown): PrimitiveType { + if (value === undefined) { + return 'undefined'; + } + + if (typeof value === 'string') { + return 'string'; + } + + if (typeof value === 'number' || typeof value === 'bigint') { + return 'number'; + } + + if (typeof value === 'boolean') { + return 'boolean'; + } + + if (value instanceof Date) { + return 'date'; + } + + if (value instanceof RegExp) { + return 'regex'; + } + + if (value instanceof Uint8Array || value instanceof Binary) { + return 'binary'; + } + + if (value instanceof Timestamp) { + return 'timestamp'; + } + + if (value instanceof ObjectId) { + return 'objectId'; + } + + if (value instanceof Decimal128 || value instanceof Double || value instanceof Int32 || value instanceof Long) { + return 'number'; + } + + const bsonType = getBsonType(value); + if (bsonType) { + return mapBsonType(bsonType); + } + + return 'unknown'; +} + +function mapBsonType(type: string): PrimitiveType { + const normalized = type.toLowerCase(); + + switch (normalized) { + case 'objectid': + return 'objectId'; + case 'decimal128': + case 'double': + case 'int32': + case 'long': + return 'number'; + case 'timestamp': + return 'timestamp'; + case 'binary': + return 'binary'; + default: + return 'unknown'; + } +} + +function getBsonType(value: unknown): string | undefined { + if (!value || typeof value !== 'object') { + return undefined; + } + + const potential = value as { _bsontype?: unknown }; + if (typeof potential._bsontype === 'string') { + return potential._bsontype; + } + + return undefined; +} + +function createFieldSummary(): FieldSummary { + return { + primitiveTypes: new Set(), + hasNull: false, + }; +} + +function createArraySummary(): ArraySummary { + return { + elementSummary: createFieldSummary(), + vectorLengths: new Set(), + nonVectorObserved: false, + sawValues: false, + }; +} + +function convertProperties(properties: Map): Record { + const result: Record = {}; + + for (const key of Array.from(properties.keys()).sort((a, b) => a.localeCompare(b))) { + result[key] = convertSummary(properties.get(key) as FieldSummary); + } + + return result; +} + +function convertSummary(summary: FieldSummary): SchemaFieldDefinition { + const variants: SchemaFieldDefinition[] = []; + + if (summary.primitiveTypes.size > 0) { + variants.push(combinePrimitiveTypes(summary.primitiveTypes)); + } + + if (summary.objectProperties && summary.objectProperties.size > 0) { + variants.push({ + type: 'object', + properties: convertProperties(summary.objectProperties), + }); + } + + if (summary.arraySummary) { + variants.push(convertArraySummary(summary.arraySummary)); + } + + if (summary.hasNull) { + variants.push('null'); + } + + if (variants.length === 0) { + return 'unknown'; + } + + if (variants.length === 1) { + return variants[0]; + } + + return { type: 'union', variants }; +} + +function combinePrimitiveTypes(types: Set): string { + const filtered = Array.from(types) + .filter((type) => type !== 'undefined') + .sort((a, b) => a.localeCompare(b)); + + if (filtered.length === 0) { + return types.has('undefined') ? 'undefined' : 'unknown'; + } + + return filtered.join(' | '); +} + +function convertArraySummary(summary: ArraySummary): SchemaArrayDefinition | SchemaUnionDefinition { + const items = convertSummary(summary.elementSummary); + const definition: SchemaArrayDefinition = { + type: 'array', + items, + }; + + if (!summary.nonVectorObserved && summary.vectorLengths.size === 1 && summary.sawValues) { + definition.vectorLength = Array.from(summary.vectorLengths)[0]; + } + + if (summary.elementSummary.hasNull && typeof items === 'string' && items === 'unknown') { + return { type: 'union', variants: [definition, 'null'] }; + } + + return definition; +} + +function isPlainObject(value: unknown): boolean { + if (typeof value !== 'object' || value === null) { + return false; + } + + if (Array.isArray(value)) { + return false; + } + + if (value instanceof Date || value instanceof RegExp) { + return false; + } + + if (value instanceof Uint8Array || value instanceof Binary) { + return false; + } + + if (getBsonType(value)) { + return false; + } + + return true; +} + +function isNumericValue(value: unknown): boolean { + if (typeof value === 'number') { + return Number.isFinite(value); + } + + if (typeof value === 'bigint') { + return true; + } + + if (value instanceof Decimal128 || value instanceof Double || value instanceof Int32 || value instanceof Long) { + return true; + } + + return getBsonType(value)?.toLowerCase() === 'decimal128'; +} diff --git a/src/webviews/documentdb/collectionView/collectionViewRouter.ts b/src/webviews/documentdb/collectionView/collectionViewRouter.ts index b60b818dd..8004236fd 100644 --- a/src/webviews/documentdb/collectionView/collectionViewRouter.ts +++ b/src/webviews/documentdb/collectionView/collectionViewRouter.ts @@ -3,6 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import { callWithTelemetryAndErrorHandling, type IActionContext } from '@microsoft/vscode-azext-utils'; import * as vscode from 'vscode'; import { type JSONSchema } from 'vscode-json-languageservice'; import { z } from 'zod'; @@ -12,6 +13,11 @@ import { getKnownFields, type FieldEntry } from '../../../utils/json/mongo/autoc import { publicProcedure, router, trpcToTelemetry } from '../../api/extension-server/trpc'; import * as l10n from '@vscode/l10n'; +import { + generateQuery, + QueryGenerationType, + type QueryGenerationContext, +} from '../../../commands/llmEnhancedCommands/queryGenerationCommands'; import { showConfirmationAsInSettings } from '../../../utils/dialogs/showConfirmation'; import { Views } from '../../../documentdb/Views'; @@ -332,54 +338,77 @@ export const collectionsViewRouter = router({ prompt: z.string(), }), ) - // procedure type - .mutation(async ({ input }) => { - await new Promise((resolve) => setTimeout(resolve, 1000)); + // handle generation request + .query(async ({ input, ctx }) => { + const generationCtx = ctx as RouterContext; + + const result = await callWithTelemetryAndErrorHandling( + 'vscode-documentdb.collectionView.generateQuery', + async (context: IActionContext) => { + // Prepare query generation context + const queryContext: QueryGenerationContext = { + clusterId: generationCtx.clusterId, + databaseName: generationCtx.databaseName, + collectionName: generationCtx.collectionName, + // For now, only handle Find queries + targetQueryType: 'Find', + naturalLanguageQuery: input.prompt, + generationType: QueryGenerationType.SingleCollection, + }; + + // Generate query with LLM + const generationResult = await generateQuery(context, queryContext); + if (generationResult.generatedQuery === undefined) { + const errorExplanation = generationResult.explanation + ? generationResult.explanation.startsWith('Error:') + ? generationResult.explanation.slice(6).trim() + : generationResult.explanation + : 'No detailed error message provided.'; + context.telemetry.properties.generationError = errorExplanation; + throw new Error(l10n.t('Query generation failed with the error: {0}', errorExplanation)); + } - // Mock AI logic: if prompt contains "error", simulate an error - if (input.prompt.toLowerCase().includes('error')) { - throw new Error('Simulated AI error for testing purposes'); - } + // Parse the generated command + // For now we only support find query + let parsedCommand: { + filter?: string; + project?: string; + sort?: string; + skip?: number; + limit?: number; + }; - // Start with current query values - const result = { - filter: input.currentQuery.filter, - project: input.currentQuery.project ?? '{ }', - sort: input.currentQuery.sort ?? '{ }', - skip: input.currentQuery.skip ?? 0, - limit: input.currentQuery.limit ?? 0, - }; - - // Mock AI logic: if prompt contains "sort" - if (input.prompt.toLowerCase().includes('sort')) { - const currentSort = input.currentQuery.sort?.trim(); - - // If there's an existing sort, reverse the directions - if (currentSort && currentSort !== '{}' && currentSort !== '{ }') { try { - // Parse the current sort to reverse field directions - const sortObj = JSON.parse(currentSort) as Record; - const reversedSort: Record = {}; - - for (const [field, direction] of Object.entries(sortObj)) { - if (typeof direction === 'number') { - // Reverse: 1 → -1, -1 → 1 - reversedSort[field] = direction === 1 ? -1 : 1; - } else { - // Keep as-is if not a number - reversedSort[field] = direction as number; - } - } - - result.sort = JSON.stringify(reversedSort, null, 0); - } catch { - // If parsing fails, use default sort - result.sort = '{ "_id": -1 }'; + parsedCommand = JSON.parse(generationResult.generatedQuery) as { + filter?: string; + project?: string; + sort?: string; + skip?: number; + limit?: number; + }; + } catch (error) { + // Add error details to telemetry + context.telemetry.properties.parseError = error instanceof Error ? error.name : 'UnknownError'; + context.telemetry.properties.parseErrorMessage = + error instanceof Error ? error.message : String(error); + + throw new Error( + l10n.t('Failed to parse generated query. Query generation provided an invalid response.'), + ); } - } else { - // No existing sort, use default - result.sort = '{ "_id": -1 }'; - } + + return { + filter: parsedCommand.filter ?? input.currentQuery.filter, + project: parsedCommand.project ?? input.currentQuery.project ?? '{ }', + sort: parsedCommand.sort ?? input.currentQuery.sort ?? '{ }', + skip: parsedCommand.skip ?? input.currentQuery.skip ?? 0, + limit: parsedCommand.limit ?? input.currentQuery.limit ?? 0, + }; + }, + ); + + if (!result) { + throw new Error(l10n.t('Query generation failed')); } return result; diff --git a/src/webviews/documentdb/collectionView/components/queryEditor/QueryEditor.tsx b/src/webviews/documentdb/collectionView/components/queryEditor/QueryEditor.tsx index f539f7f05..e69fcc7ba 100644 --- a/src/webviews/documentdb/collectionView/components/queryEditor/QueryEditor.tsx +++ b/src/webviews/documentdb/collectionView/components/queryEditor/QueryEditor.tsx @@ -44,6 +44,7 @@ export const QueryEditor = ({ onExecuteRequest }: QueryEditorProps): JSX.Element const [aiPromptHistory, setAiPromptHistory] = useState([]); const schemaAbortControllerRef = useRef(null); + const aiGenerationAbortControllerRef = useRef(null); const aiInputRef = useRef(null); // Refs for Monaco editors @@ -176,6 +177,10 @@ export const QueryEditor = ({ onExecuteRequest }: QueryEditorProps): JSX.Element schemaAbortControllerRef.current.abort(); schemaAbortControllerRef.current = null; } + if (aiGenerationAbortControllerRef.current) { + aiGenerationAbortControllerRef.current.abort(); + aiGenerationAbortControllerRef.current = null; + } }; }, []); @@ -253,10 +258,19 @@ export const QueryEditor = ({ onExecuteRequest }: QueryEditorProps): JSX.Element return; // Don't generate if prompt is empty } + // Cancel any previous AI generation request by marking it as aborted + if (aiGenerationAbortControllerRef.current) { + aiGenerationAbortControllerRef.current.abort(); + } + + // Create new AbortController for this request (used for client-side cancellation only) + const abortController = new AbortController(); + aiGenerationAbortControllerRef.current = abortController; + setIsAiActive(true); try { - const result = await trpcClient.mongoClusters.collectionView.generateQuery.mutate({ + const result = await trpcClient.mongoClusters.collectionView.generateQuery.query({ currentQuery: { filter: filterValue, project: projectValue, @@ -267,6 +281,11 @@ export const QueryEditor = ({ onExecuteRequest }: QueryEditorProps): JSX.Element prompt: aiPromptValue, }); + // Check if this request was aborted while waiting for response + if (abortController.signal.aborted) { + return; // Ignore the response if we aborted + } + // Update state with generated query setFilterValue(result.filter); setProjectValue(result.project); @@ -285,13 +304,22 @@ export const QueryEditor = ({ onExecuteRequest }: QueryEditorProps): JSX.Element // Clear the AI prompt after successful generation setAiPromptValue(''); } catch (error) { + // Check if this request was aborted + if (abortController.signal.aborted) { + return; // Ignore errors from aborted requests + } + void trpcClient.common.displayErrorMessage.mutate({ message: l10n.t('Error generating query'), modal: false, cause: error instanceof Error ? error.message : String(error), }); } finally { - setIsAiActive(false); + // Only clear active state if this request wasn't aborted + if (!abortController.signal.aborted) { + setIsAiActive(false); + aiGenerationAbortControllerRef.current = null; + } } }; @@ -334,15 +362,26 @@ export const QueryEditor = ({ onExecuteRequest }: QueryEditorProps): JSX.Element // Execute query on Ctrl+Enter / Cmd+Enter onExecuteRequest(); } else if (event.key === 'Escape') { - // ESC key - hide AI area and reset active state - setIsAiActive(false); - setAiPromptValue(''); // Clear the prompt - setCurrentContext((prev) => ({ - ...prev, - isAiRowVisible: false, - })); - // Give focus to the filter editor - filterEditorRef.current?.focus(); + // ESC key behavior: + // - If AI is active (generation in progress): cancel the request + // - If AI is not active: hide AI row and clear prompt + if (isAiActive) { + // Cancel the ongoing AI generation + if (aiGenerationAbortControllerRef.current) { + aiGenerationAbortControllerRef.current.abort(); + aiGenerationAbortControllerRef.current = null; + } + setIsAiActive(false); + } else { + // Hide AI area and clear prompt + setAiPromptValue(''); + setCurrentContext((prev) => ({ + ...prev, + isAiRowVisible: false, + })); + // Give focus to the filter editor + filterEditorRef.current?.focus(); + } } }} />