diff --git a/src/package.json b/src/package.json index 34808a28c0..cb799920a9 100644 --- a/src/package.json +++ b/src/package.json @@ -407,6 +407,78 @@ "minimum": 1, "maximum": 200, "description": "%settings.codeIndex.embeddingBatchSize.description%" + }, + "roo-cline.subLlm.enabled": { + "type": "boolean", + "default": false, + "description": "Enable sub-LLM features for code search enhancement" + }, + "roo-cline.subLlm.model.mode": { + "type": "string", + "enum": [ + "mirror", + "custom" + ], + "default": "mirror", + "description": "Model selection mode: 'mirror' uses chat model, 'custom' allows override" + }, + "roo-cline.subLlm.model.provider": { + "type": "string", + "default": "", + "description": "Custom provider for sub-LLM when mode is 'custom'" + }, + "roo-cline.subLlm.model.modelId": { + "type": "string", + "default": "", + "description": "Custom model ID for sub-LLM when mode is 'custom'" + }, + "roo-cline.subLlm.maxTokensPerOp": { + "type": "number", + "default": 1000, + "minimum": 100, + "maximum": 4000, + "description": "Maximum tokens per sub-LLM operation" + }, + "roo-cline.subLlm.dailyCostCapUSD": { + "type": "number", + "default": 1, + "minimum": 0, + "maximum": 100, + "description": "Daily cost cap in USD for sub-LLM operations" + }, + "roo-cline.subLlm.timeout": { + "type": "number", + "default": 5000, + "minimum": 1000, + "maximum": 30000, + "description": "Timeout in milliseconds for sub-LLM operations" + }, + "roo-cline.codeIndex.llm.rewriter": { + "type": "boolean", + "default": false, + "description": "Enable query rewriting for natural language search" + }, + "roo-cline.codeIndex.llm.reranker": { + "type": "boolean", + "default": false, + "description": "Enable LLM-based reranking of search results" + }, + "roo-cline.codeIndex.llm.summaries": { + "type": "boolean", + "default": false, + "description": "Enable generation of code summaries (requires reindex)" + }, + "roo-cline.codeIndex.llm.tags": { + "type": "boolean", + "default": false, + "description": "Enable generation of code tags (requires reindex)" + }, + "roo-cline.codeIndex.llm.maxKForRerank": { + "type": "number", + "default": 50, + "minimum": 10, + "maximum": 200, + "description": "Maximum number of candidates to rerank" } } } diff --git a/src/services/code-index/enhanced-search-service.ts b/src/services/code-index/enhanced-search-service.ts new file mode 100644 index 0000000000..4f813bd880 --- /dev/null +++ b/src/services/code-index/enhanced-search-service.ts @@ -0,0 +1,246 @@ +import * as path from "path" +import * as vscode from "vscode" +import { VectorStoreSearchResult } from "./interfaces" +import { IEmbedder } from "./interfaces/embedder" +import { IVectorStore } from "./interfaces/vector-store" +import { CodeIndexConfigManager } from "./config-manager" +import { CodeIndexStateManager } from "./state-manager" +import { TelemetryService } from "@roo-code/telemetry" +import { TelemetryEventName } from "@roo-code/types" +import { LlmClient, QueryRewriter, Reranker, SubLlmConfig, QueryVariant } from "../llm-utils" +import { Package } from "../../shared/package" + +/** + * Enhanced search service with LLM-assisted query rewriting and reranking + */ +export class EnhancedCodeIndexSearchService { + private llmClient: LlmClient | null = null + private queryRewriter: QueryRewriter | null = null + private reranker: Reranker | null = null + private subLlmConfig: SubLlmConfig | null = null + + constructor( + private readonly configManager: CodeIndexConfigManager, + private readonly stateManager: CodeIndexStateManager, + private readonly embedder: IEmbedder, + private readonly vectorStore: IVectorStore, + private readonly context: vscode.ExtensionContext, + ) { + this.initializeLlmComponents() + } + + /** + * Initialize LLM components if enabled + */ + private async initializeLlmComponents(): Promise { + try { + // Get sub-LLM configuration from VSCode settings + const config = vscode.workspace.getConfiguration(Package.name) + + this.subLlmConfig = { + enabled: config.get("subLlm.enabled", false), + modelMode: config.get<"mirror" | "custom">("subLlm.model.mode", "mirror"), + maxTokensPerOp: config.get("subLlm.maxTokensPerOp", 1000), + dailyCostCapUSD: config.get("subLlm.dailyCostCapUSD", 1.0), + timeout: config.get("subLlm.timeout", 5000), + } + + if (this.subLlmConfig.enabled) { + // Initialize LLM client + this.llmClient = new LlmClient(this.context, this.subLlmConfig) + await this.llmClient.initialize() + + // Initialize components + this.queryRewriter = new QueryRewriter(this.llmClient) + this.reranker = new Reranker(this.llmClient) + } + } catch (error) { + console.error("[EnhancedSearchService] Failed to initialize LLM components:", error) + // Continue without LLM features + } + } + + /** + * Enhanced search with optional query rewriting and reranking + */ + public async searchIndex(query: string, directoryPrefix?: string): Promise { + if (!this.configManager.isFeatureEnabled || !this.configManager.isFeatureConfigured) { + throw new Error("Code index feature is disabled or not configured.") + } + + const minScore = this.configManager.currentSearchMinScore + const maxResults = this.configManager.currentSearchMaxResults + + const currentState = this.stateManager.getCurrentStatus().systemStatus + if (currentState !== "Indexed" && currentState !== "Indexing") { + throw new Error(`Code index is not ready for search. Current state: ${currentState}`) + } + + try { + // Step 1: Query Rewriting (if enabled) + const queries = await this.getSearchQueries(query) + + // Step 2: Perform searches for all query variants + const allResults = await this.performMultiSearch(queries, directoryPrefix, minScore, maxResults) + + // Step 3: Deduplicate and merge results + const mergedResults = this.mergeSearchResults(allResults) + + // Step 4: Rerank results (if enabled) + const finalResults = await this.rerankResults(query, mergedResults) + + // Step 5: Apply final filtering and limit + return this.filterAndLimitResults(finalResults, minScore, maxResults) + } catch (error) { + console.error("[EnhancedSearchService] Error during search:", error) + this.stateManager.setSystemState("Error", `Search failed: ${(error as Error).message}`) + + TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { + error: (error as Error).message, + stack: (error as Error).stack, + location: "enhancedSearchIndex", + }) + + throw error + } + } + + /** + * Get search queries (original + rewritten variants if enabled) + */ + private async getSearchQueries(query: string): Promise { + const queries = [query] // Always include original + + // Check if query rewriting is enabled + const rewriterEnabled = vscode.workspace + .getConfiguration(Package.name) + .get("codeIndex.llm.rewriter", false) + + if (rewriterEnabled && this.queryRewriter) { + try { + const variants = await this.queryRewriter.rewrite(query) + // Add variant queries (excluding the original which is already included) + variants + .filter((v: QueryVariant) => v.type !== "original") + .forEach((v: QueryVariant) => queries.push(v.query)) + } catch (error) { + console.warn("[EnhancedSearchService] Query rewriting failed, using original only:", error) + } + } + + return queries + } + + /** + * Perform searches for multiple query variants + */ + private async performMultiSearch( + queries: string[], + directoryPrefix: string | undefined, + minScore: number, + maxResults: number, + ): Promise { + const normalizedPrefix = directoryPrefix ? path.normalize(directoryPrefix) : undefined + + const searchPromises = queries.map(async (q) => { + try { + // Generate embedding for query + const embeddingResponse = await this.embedder.createEmbeddings([q]) + const vector = embeddingResponse?.embeddings[0] + + if (!vector) { + console.warn(`Failed to generate embedding for query variant: ${q}`) + return [] + } + + // Perform search + return await this.vectorStore.search(vector, normalizedPrefix, minScore, maxResults) + } catch (error) { + console.warn(`Search failed for query variant: ${q}`, error) + return [] + } + }) + + return Promise.all(searchPromises) + } + + /** + * Merge and deduplicate search results from multiple queries + */ + private mergeSearchResults(allResults: VectorStoreSearchResult[][]): VectorStoreSearchResult[] { + const resultMap = new Map() + + // Merge results, keeping the highest score for duplicates + for (const results of allResults) { + for (const result of results) { + const existing = resultMap.get(result.id) + if (!existing || result.score > existing.score) { + resultMap.set(result.id, result) + } + } + } + + // Convert back to array and sort by score + return Array.from(resultMap.values()).sort((a, b) => b.score - a.score) + } + + /** + * Rerank results using LLM if enabled + */ + private async rerankResults(query: string, results: VectorStoreSearchResult[]): Promise { + // Check if reranking is enabled + const rerankerEnabled = vscode.workspace + .getConfiguration(Package.name) + .get("codeIndex.llm.reranker", false) + + if (!rerankerEnabled || !this.reranker || results.length === 0) { + return results + } + + try { + // Get max candidates for reranking + const maxKForRerank = vscode.workspace + .getConfiguration(Package.name) + .get("codeIndex.llm.maxKForRerank", 50) + + // Rerank top-K candidates + const rerankResults = await this.reranker.rerank(query, results, { + maxCandidates: maxKForRerank, + includeReason: true, + }) + + // Blend reranked scores with original scores + return this.reranker.blendScores(rerankResults, results, 0.7) + } catch (error) { + console.warn("[EnhancedSearchService] Reranking failed, using original scores:", error) + return results + } + } + + /** + * Apply final filtering and limit to results + */ + private filterAndLimitResults( + results: VectorStoreSearchResult[], + minScore: number, + maxResults: number, + ): VectorStoreSearchResult[] { + return results.filter((r) => r.score >= minScore).slice(0, maxResults) + } + + /** + * Get LLM budget status + */ + public getLlmBudgetStatus(): { enabled: boolean; dailyCost?: number; remaining?: number } { + if (!this.llmClient) { + return { enabled: false } + } + + const status = this.llmClient.getBudgetStatus() + return { + enabled: true, + dailyCost: status.dailyCost, + remaining: status.remaining || undefined, + } + } +} diff --git a/src/services/code-index/processors/enhanced-scanner.ts b/src/services/code-index/processors/enhanced-scanner.ts new file mode 100644 index 0000000000..4c038e99d5 --- /dev/null +++ b/src/services/code-index/processors/enhanced-scanner.ts @@ -0,0 +1,185 @@ +import { Ignore } from "ignore" +import * as vscode from "vscode" +import { DirectoryScanner } from "./scanner" +import { ICodeParser, IEmbedder, IVectorStore, IDirectoryScanner, CodeBlock } from "../interfaces" +import { CacheManager } from "../cache-manager" +import { LlmClient, Summarizer, SubLlmConfig } from "../../llm-utils" +import { Package } from "../../../shared/package" + +/** + * Factory for creating enhanced directory scanners with optional LLM features + */ +export class EnhancedScannerFactory { + /** + * Create a directory scanner with optional LLM-based summarization + */ + static async create( + embedder: IEmbedder, + qdrantClient: IVectorStore, + codeParser: ICodeParser, + cacheManager: CacheManager, + ignoreInstance: Ignore, + context: vscode.ExtensionContext, + batchSegmentThreshold?: number, + ): Promise { + const config = vscode.workspace.getConfiguration(Package.name) + + // Check if LLM features are enabled + const summariesEnabled = config.get("codeIndex.llm.summaries", false) + const tagsEnabled = config.get("codeIndex.llm.tags", false) + + if (summariesEnabled || tagsEnabled) { + // Create enhanced scanner with LLM features + return new EnhancedDirectoryScanner( + embedder, + qdrantClient, + codeParser, + cacheManager, + ignoreInstance, + context, + summariesEnabled, + tagsEnabled, + batchSegmentThreshold, + ) + } + + // Return standard scanner + return new DirectoryScanner( + embedder, + qdrantClient, + codeParser, + cacheManager, + ignoreInstance, + batchSegmentThreshold, + ) + } +} + +/** + * Enhanced directory scanner with LLM-based summarization + * This is a wrapper around DirectoryScanner that adds summarization capabilities + */ +class EnhancedDirectoryScanner implements IDirectoryScanner { + private scanner: DirectoryScanner + private llmClient: LlmClient | null = null + private summarizer: Summarizer | null = null + + constructor( + private readonly embedder: IEmbedder, + private readonly qdrantClient: IVectorStore, + private readonly codeParser: ICodeParser, + private readonly cacheManager: CacheManager, + private readonly ignoreInstance: Ignore, + private readonly context: vscode.ExtensionContext, + private readonly summariesEnabled: boolean, + private readonly tagsEnabled: boolean, + batchSegmentThreshold?: number, + ) { + // Create wrapped scanner + this.scanner = new DirectoryScanner( + new EnhancedEmbedder(embedder, this), + qdrantClient, + codeParser, + cacheManager, + ignoreInstance, + batchSegmentThreshold, + ) + + this.initializeLlmComponents() + } + + /** + * Initialize LLM components + */ + private async initializeLlmComponents(): Promise { + try { + const config = vscode.workspace.getConfiguration(Package.name) + + // Initialize sub-LLM configuration + const subLlmConfig: SubLlmConfig = { + enabled: config.get("subLlm.enabled", false), + modelMode: config.get<"mirror" | "custom">("subLlm.model.mode", "mirror"), + maxTokensPerOp: config.get("subLlm.maxTokensPerOp", 500), + dailyCostCapUSD: config.get("subLlm.dailyCostCapUSD", 1.0), + timeout: config.get("subLlm.timeout", 5000), + } + + if (subLlmConfig.enabled) { + this.llmClient = new LlmClient(this.context, subLlmConfig) + await this.llmClient.initialize() + this.summarizer = new Summarizer(this.llmClient) + } + } catch (error) { + console.error("[EnhancedScanner] Failed to initialize LLM components:", error) + } + } + + /** + * Scan directory with enhanced features + */ + async scanDirectory( + directory: string, + onError?: (error: Error) => void, + onBlocksIndexed?: (indexedCount: number) => void, + onFileParsed?: (fileBlockCount: number) => void, + ): Promise<{ + stats: { processed: number; skipped: number } + totalBlockCount: number + }> { + return this.scanner.scanDirectory(directory, onError, onBlocksIndexed, onFileParsed) + } + + /** + * Get summarizer for external use + */ + getSummarizer(): Summarizer | null { + return this.summarizer + } + + /** + * Check if summaries are enabled + */ + isSummariesEnabled(): boolean { + return this.summariesEnabled + } + + /** + * Check if tags are enabled + */ + isTagsEnabled(): boolean { + return this.tagsEnabled + } +} + +/** + * Enhanced embedder that augments text with summaries + */ +class EnhancedEmbedder implements IEmbedder { + constructor( + private readonly baseEmbedder: IEmbedder, + private readonly enhancedScanner: EnhancedDirectoryScanner, + ) {} + + get embedderInfo() { + return this.baseEmbedder.embedderInfo + } + + async createEmbeddings(texts: string[], model?: string) { + // Check if we should augment with summaries + const summarizer = this.enhancedScanner.getSummarizer() + + if (!summarizer || (!this.enhancedScanner.isSummariesEnabled() && !this.enhancedScanner.isTagsEnabled())) { + // No augmentation needed + return this.baseEmbedder.createEmbeddings(texts, model) + } + + // For now, pass through to base embedder + // In a full implementation, we would parse the text to extract code blocks, + // generate summaries, and augment the text before embedding + return this.baseEmbedder.createEmbeddings(texts, model) + } + + async validateConfiguration() { + return this.baseEmbedder.validateConfiguration() + } +} diff --git a/src/services/llm-utils/index.ts b/src/services/llm-utils/index.ts new file mode 100644 index 0000000000..5ba85e85c6 --- /dev/null +++ b/src/services/llm-utils/index.ts @@ -0,0 +1,6 @@ +export { LlmClient } from "./llm-client" +export { JsonRunner } from "./json-runner" +export { QueryRewriter } from "./query-rewriter" +export { Reranker } from "./reranker" +export { Summarizer } from "./summarizer" +export type { SubLlmConfig, SubLlmMode, QueryVariant, RerankResult, SummaryResult } from "./types" diff --git a/src/services/llm-utils/json-runner.ts b/src/services/llm-utils/json-runner.ts new file mode 100644 index 0000000000..44c938edc6 --- /dev/null +++ b/src/services/llm-utils/json-runner.ts @@ -0,0 +1,199 @@ +import { z } from "zod" + +/** + * Handles JSON extraction, validation, and sanitization for LLM responses + */ +export class JsonRunner { + /** + * Extract and validate JSON from LLM response + */ + static extract(text: string, schema?: z.ZodSchema): T { + // Strip markdown fences + const stripped = this.stripFences(text) + + // Find first JSON object/array + const json = this.findFirstJson(stripped) + + if (!json) { + throw new Error("No valid JSON found in response") + } + + // Parse JSON + let parsed: any + try { + parsed = JSON.parse(json) + } catch (error) { + throw new Error(`Invalid JSON: ${error instanceof Error ? error.message : String(error)}`) + } + + // Validate with schema if provided + if (schema) { + const result = schema.safeParse(parsed) + if (!result.success) { + throw new Error(`Schema validation failed: ${result.error.message}`) + } + return result.data + } + + return parsed as T + } + + /** + * Strip markdown code fences from text + */ + private static stripFences(text: string): string { + // Remove ```json ... ``` or ``` ... ``` + const fencePattern = /```(?:json)?\s*([\s\S]*?)```/g + const matches = [...text.matchAll(fencePattern)] + + if (matches.length > 0) { + // Return content from first fence + return matches[0][1].trim() + } + + return text.trim() + } + + /** + * Find first complete JSON object or array in text + */ + private static findFirstJson(text: string): string | null { + const trimmed = text.trim() + + // Try to find JSON object or array + const objectStart = trimmed.indexOf("{") + const arrayStart = trimmed.indexOf("[") + + if (objectStart === -1 && arrayStart === -1) { + // No JSON markers found + return null + } + + // Determine which comes first + let start: number + let openChar: string + let closeChar: string + + if (objectStart !== -1 && (arrayStart === -1 || objectStart < arrayStart)) { + start = objectStart + openChar = "{" + closeChar = "}" + } else { + start = arrayStart + openChar = "[" + closeChar = "]" + } + + // Extract JSON using bracket counting + return this.extractJsonByBrackets(trimmed, start, openChar, closeChar) + } + + /** + * Extract JSON by counting brackets + */ + private static extractJsonByBrackets( + text: string, + start: number, + openChar: string, + closeChar: string, + ): string | null { + let depth = 0 + let inString = false + let escape = false + + for (let i = start; i < text.length; i++) { + const char = text[i] + + // Handle escape sequences + if (escape) { + escape = false + continue + } + + if (char === "\\") { + escape = true + continue + } + + // Handle strings + if (char === '"' && !escape) { + inString = !inString + continue + } + + // Count brackets only outside strings + if (!inString) { + if (char === openChar) { + depth++ + } else if (char === closeChar) { + depth-- + if (depth === 0) { + return text.substring(start, i + 1) + } + } + } + } + + // No complete JSON found + return null + } + + /** + * Sanitize JSON for logging (remove sensitive data) + */ + static sanitize(obj: any, sensitiveKeys: string[] = []): any { + const defaultSensitiveKeys = ["apiKey", "api_key", "token", "secret", "password", "credential", "auth"] + + const allSensitiveKeys = [...defaultSensitiveKeys, ...sensitiveKeys] + + if (obj === null || obj === undefined) { + return obj + } + + if (typeof obj !== "object") { + return obj + } + + if (Array.isArray(obj)) { + return obj.map((item) => this.sanitize(item, sensitiveKeys)) + } + + const sanitized: any = {} + for (const [key, value] of Object.entries(obj)) { + const lowerKey = key.toLowerCase() + const isSensitive = allSensitiveKeys.some((sensitive) => lowerKey.includes(sensitive.toLowerCase())) + + if (isSensitive) { + sanitized[key] = "[REDACTED]" + } else if (typeof value === "object" && value !== null) { + sanitized[key] = this.sanitize(value, sensitiveKeys) + } else { + sanitized[key] = value + } + } + + return sanitized + } + + /** + * Validate JSON against a Zod schema with detailed error reporting + */ + static validate( + data: unknown, + schema: z.ZodSchema, + ): { success: true; data: T } | { success: false; errors: string[] } { + const result = schema.safeParse(data) + + if (result.success) { + return { success: true, data: result.data } + } + + // Extract detailed error messages + const errors = result.error.errors.map((err) => { + const path = err.path.join(".") + return path ? `${path}: ${err.message}` : err.message + }) + + return { success: false, errors } + } +} diff --git a/src/services/llm-utils/llm-client.ts b/src/services/llm-utils/llm-client.ts new file mode 100644 index 0000000000..42b8cbdfe3 --- /dev/null +++ b/src/services/llm-utils/llm-client.ts @@ -0,0 +1,254 @@ +import { Anthropic } from "@anthropic-ai/sdk" +import { ProviderSettings } from "@roo-code/types" +import { buildApiHandler, ApiHandler } from "../../api" +import { SubLlmConfig, SubLlmMode, GenerateOptions, GenerateJsonOptions } from "./types" +import { ProviderSettingsManager } from "../../core/config/ProviderSettingsManager" +import * as vscode from "vscode" + +/** + * Client for sub-LLM operations with provider selection and budget management + */ +export class LlmClient { + private apiHandler: ApiHandler | null = null + private config: SubLlmConfig + private providerSettings: ProviderSettings | null = null + private dailyCost: number = 0 + private lastCostReset: Date = new Date() + + constructor( + private readonly context: vscode.ExtensionContext, + config: SubLlmConfig, + ) { + this.config = config + } + + /** + * Initialize the LLM client with appropriate provider + */ + async initialize(): Promise { + if (!this.config.enabled) { + return + } + + // Get provider settings based on mode + if (this.config.modelMode === "mirror") { + // Mirror the chat model settings + const settingsManager = new ProviderSettingsManager(this.context) + const profiles = await settingsManager.export() + const currentProfile = profiles.apiConfigs[profiles.currentApiConfigName] + + if (currentProfile) { + this.providerSettings = currentProfile as ProviderSettings + } + } else if (this.config.modelMode === "custom" && this.config.customProvider) { + // Use custom provider settings + this.providerSettings = this.config.customProvider + } + + if (this.providerSettings) { + this.apiHandler = buildApiHandler(this.providerSettings) + } + } + + /** + * Check and update daily cost budget + */ + private checkBudget(estimatedCost: number): boolean { + const now = new Date() + + // Reset daily cost if it's a new day + if (now.toDateString() !== this.lastCostReset.toDateString()) { + this.dailyCost = 0 + this.lastCostReset = now + } + + // Check if adding this cost would exceed the cap + if (this.config.dailyCostCapUSD && this.dailyCost + estimatedCost > this.config.dailyCostCapUSD) { + return false + } + + return true + } + + /** + * Update the daily cost tracker + */ + private updateCost(cost: number): void { + this.dailyCost += cost + } + + /** + * Generate text completion + */ + async generateText(prompt: string, options: GenerateOptions = {}): Promise { + if (!this.config.enabled || !this.apiHandler) { + throw new Error("LLM client not initialized or disabled") + } + + // Estimate cost (simplified - would need actual token counting) + const estimatedCost = 0.001 // Placeholder + if (!this.checkBudget(estimatedCost)) { + throw new Error("Daily LLM cost budget exceeded") + } + + const systemPrompt = options.systemPrompt || "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: prompt, + }, + ] + + try { + const stream = this.apiHandler.createMessage(systemPrompt, messages) + let result = "" + + for await (const chunk of stream) { + if (chunk.type === "text") { + result += chunk.text + } else if (chunk.type === "usage") { + // Calculate actual cost based on usage + const actualCost = this.calculateCost(chunk.inputTokens || 0, chunk.outputTokens || 0) + this.updateCost(actualCost) + } + } + + return result + } catch (error) { + console.error("[LlmClient] Error generating text:", error) + throw error + } + } + + /** + * Generate JSON with validation + */ + async generateJson(prompt: string, options: GenerateJsonOptions = {}): Promise { + const jsonPrompt = `${prompt}\n\nRespond with valid JSON only, no additional text or markdown.` + + let attempts = 0 + const maxRetries = options.maxRetries || 3 + + while (attempts < maxRetries) { + attempts++ + + try { + const response = await this.generateText(jsonPrompt, options) + + // Extract JSON from response (handle markdown fences) + const jsonStr = this.extractJson(response) + const parsed = JSON.parse(jsonStr) + + // Validate with schema if provided + if (options.schema) { + const result = options.schema.safeParse(parsed) + if (!result.success) { + if (options.retryOnValidationFailure && attempts < maxRetries) { + continue + } + throw new Error(`JSON validation failed: ${result.error.message}`) + } + return result.data + } + + return parsed as T + } catch (error) { + if (attempts >= maxRetries) { + throw error + } + } + } + + throw new Error("Failed to generate valid JSON after retries") + } + + /** + * Extract JSON from a string that might contain markdown fences + */ + private extractJson(text: string): string { + // Remove markdown code fences + const fencePattern = /```(?:json)?\s*([\s\S]*?)```/ + const match = text.match(fencePattern) + if (match) { + return match[1].trim() + } + + // Try to find first complete JSON object + const firstBrace = text.indexOf("{") + const firstBracket = text.indexOf("[") + + if (firstBrace === -1 && firstBracket === -1) { + return text.trim() + } + + const start = + firstBrace !== -1 && (firstBracket === -1 || firstBrace < firstBracket) ? firstBrace : firstBracket + + // Simple bracket counting to find end + let depth = 0 + let inString = false + let escape = false + + for (let i = start; i < text.length; i++) { + const char = text[i] + + if (escape) { + escape = false + continue + } + + if (char === "\\") { + escape = true + continue + } + + if (char === '"') { + inString = !inString + continue + } + + if (!inString) { + if (char === "{" || char === "[") { + depth++ + } else if (char === "}" || char === "]") { + depth-- + if (depth === 0) { + return text.substring(start, i + 1) + } + } + } + } + + return text.substring(start).trim() + } + + /** + * Calculate cost based on token usage + * This is a simplified version - actual costs vary by model + */ + private calculateCost(inputTokens: number, outputTokens: number): number { + // Simplified cost calculation (would need model-specific rates) + const inputCostPer1k = 0.003 + const outputCostPer1k = 0.015 + + return (inputTokens / 1000) * inputCostPer1k + (outputTokens / 1000) * outputCostPer1k + } + + /** + * Get current budget status + */ + getBudgetStatus(): { dailyCost: number; remaining: number | null } { + const now = new Date() + + // Reset if new day + if (now.toDateString() !== this.lastCostReset.toDateString()) { + this.dailyCost = 0 + this.lastCostReset = now + } + + return { + dailyCost: this.dailyCost, + remaining: this.config.dailyCostCapUSD ? this.config.dailyCostCapUSD - this.dailyCost : null, + } + } +} diff --git a/src/services/llm-utils/query-rewriter.ts b/src/services/llm-utils/query-rewriter.ts new file mode 100644 index 0000000000..a9b346fedf --- /dev/null +++ b/src/services/llm-utils/query-rewriter.ts @@ -0,0 +1,113 @@ +import { z } from "zod" +import { LlmClient } from "./llm-client" +import { JsonRunner } from "./json-runner" +import { QueryVariant } from "./types" + +// Schema for query rewriting response +const queryRewriteSchema = z.object({ + variants: z.array( + z.object({ + query: z.string(), + type: z.enum(["original", "synonym", "symbol", "natural", "error_signature"]), + reason: z.string().optional(), + }), + ), +}) + +/** + * Rewrites natural language queries into multiple code-aware variants + */ +export class QueryRewriter { + constructor(private readonly llmClient: LlmClient) {} + + /** + * Expand a single query into multiple variants for better search coverage + */ + async rewrite(query: string, context?: RewriteContext): Promise { + const systemPrompt = `You are a code search query optimizer. Your task is to rewrite natural language queries into multiple variants that will help find relevant code. + +Consider: +- Synonyms and alternative phrasings +- Technical terms vs natural language +- Symbol names and identifiers +- Common code patterns +- Error messages and signatures + +Always include the original query as the first variant. +Limit to 3 variants total. +Respond with JSON only.` + + const userPrompt = this.buildPrompt(query, context) + + try { + const response = await this.llmClient.generateJson>(userPrompt, { + schema: queryRewriteSchema, + systemPrompt, + maxTokens: 500, + temperature: 0.3, + }) + + return response.variants + } catch (error) { + console.error("[QueryRewriter] Failed to rewrite query:", error) + // Fallback to original query only + return [{ query, type: "original" }] + } + } + + /** + * Build the prompt for query rewriting + */ + private buildPrompt(query: string, context?: RewriteContext): string { + let prompt = `Rewrite this code search query into variants: +Query: "${query}" + +${context?.language ? `Language: ${context.language}` : ""} +${context?.framework ? `Framework: ${context.framework}` : ""} + +Example response format: +{ + "variants": [ + {"query": "${query}", "type": "original"}, + {"query": "alternative query 1", "type": "synonym", "reason": "uses technical terms"}, + {"query": "alternative query 2", "type": "symbol", "reason": "likely function name"} + ] +}` + + return prompt + } + + /** + * Batch rewrite multiple queries + */ + async rewriteBatch(queries: string[], context?: RewriteContext): Promise> { + const results = new Map() + + // Process in parallel with concurrency limit + const batchSize = 3 + for (let i = 0; i < queries.length; i += batchSize) { + const batch = queries.slice(i, i + batchSize) + const batchResults = await Promise.all(batch.map((q) => this.rewrite(q, context))) + + batch.forEach((query, index) => { + results.set(query, batchResults[index]) + }) + } + + return results + } +} + +/** + * Context for query rewriting + */ +export interface RewriteContext { + /** Programming language context */ + language?: string + /** Framework or library context */ + framework?: string + /** File types to search */ + fileTypes?: string[] + /** Additional context about the codebase */ + additionalContext?: string +} diff --git a/src/services/llm-utils/reranker.ts b/src/services/llm-utils/reranker.ts new file mode 100644 index 0000000000..b8419b72cc --- /dev/null +++ b/src/services/llm-utils/reranker.ts @@ -0,0 +1,190 @@ +import { z } from "zod" +import { LlmClient } from "./llm-client" +import { RerankResult } from "./types" +import { VectorStoreSearchResult } from "../code-index/interfaces" + +// Schema for reranking response +const rerankResponseSchema = z.object({ + results: z.array( + z.object({ + id: z.union([z.string(), z.number()]), + score: z.number().min(0).max(1), + reason: z.string().optional(), + }), + ), +}) + +/** + * LLM-based reranker for search results + */ +export class Reranker { + constructor(private readonly llmClient: LlmClient) {} + + /** + * Rerank search results based on relevance to query + */ + async rerank( + query: string, + candidates: VectorStoreSearchResult[], + options: RerankOptions = {}, + ): Promise { + if (candidates.length === 0) { + return [] + } + + const maxCandidates = options.maxCandidates || 50 + const topCandidates = candidates.slice(0, maxCandidates) + + const systemPrompt = `You are a code search relevance expert. Score each code snippet's relevance to the query. + +Scoring criteria: +- 1.0: Perfect match - exactly what the query is looking for +- 0.8-0.9: Highly relevant - directly addresses the query +- 0.6-0.7: Relevant - related to the query but not exact +- 0.4-0.5: Somewhat relevant - tangentially related +- 0.2-0.3: Minimally relevant - only loosely connected +- 0.0-0.1: Not relevant + +Consider: +- Semantic similarity to the query intent +- Code functionality and purpose +- Symbol names and identifiers +- Comments and documentation +- Error messages or patterns + +Respond with JSON only.` + + const userPrompt = this.buildPrompt(query, topCandidates, options) + + try { + const response = await this.llmClient.generateJson>(userPrompt, { + schema: rerankResponseSchema, + systemPrompt, + maxTokens: 1000, + temperature: 0.1, + }) + + return response.results + } catch (error) { + console.error("[Reranker] Failed to rerank results:", error) + // Fallback to original scores + return topCandidates.map((candidate, index) => ({ + id: candidate.id, + score: candidate.score, + reason: "Fallback to embedding score", + })) + } + } + + /** + * Build the prompt for reranking + */ + private buildPrompt(query: string, candidates: VectorStoreSearchResult[], options: RerankOptions): string { + const candidateTexts = candidates + .map((c, i) => { + const payload = c.payload + if (!payload) return `[${i}]: No content` + + // Truncate code to reasonable length + const code = this.truncateCode(payload.codeChunk || "", options.maxCodeLength || 350) + + return `[${c.id}]: +File: ${payload.filePath || "unknown"} +Lines: ${payload.startLine || 0}-${payload.endLine || 0} +Code: +${code} +---` + }) + .join("\n\n") + + return `Query: "${query}" + +Candidates to score: + +${candidateTexts} + +Score each candidate's relevance to the query. Return JSON with format: +{ + "results": [ + {"id": "candidate_id", "score": 0.95, "reason": "Contains exact function"}, + ... + ] +}` + } + + /** + * Truncate code snippet to maximum length while preserving structure + */ + private truncateCode(code: string, maxLength: number): string { + if (code.length <= maxLength) { + return code + } + + // Try to truncate at a natural boundary + const truncated = code.substring(0, maxLength) + const lastNewline = truncated.lastIndexOf("\n") + + if (lastNewline > maxLength * 0.8) { + return truncated.substring(0, lastNewline) + "\n..." + } + + return truncated + "..." + } + + /** + * Blend reranked scores with original embedding scores + */ + blendScores( + rerankResults: RerankResult[], + originalResults: VectorStoreSearchResult[], + blendWeight: number = 0.7, + ): VectorStoreSearchResult[] { + // Create a map of rerank scores + const rerankMap = new Map() + rerankResults.forEach((r) => rerankMap.set(r.id, r)) + + // Blend scores + const blended = originalResults.map((original) => { + const reranked = rerankMap.get(original.id) + + if (!reranked) { + return original + } + + // Blend: weight * rerank + (1-weight) * embedding + const blendedScore = blendWeight * reranked.score + (1 - blendWeight) * original.score + + return { + ...original, + score: blendedScore, + // Add rerank reason to payload if available + payload: original.payload + ? { + ...original.payload, + rerankReason: reranked.reason, + } + : undefined, + } as VectorStoreSearchResult + }) + + // Sort by blended score + return blended.sort((a, b) => b.score - a.score) + } +} + +/** + * Options for reranking + */ +export interface RerankOptions { + /** Maximum number of candidates to rerank */ + maxCandidates?: number + /** Maximum code length per candidate */ + maxCodeLength?: number + /** Include reasoning in results */ + includeReason?: boolean + /** Additional context for reranking */ + context?: { + language?: string + framework?: string + } +} diff --git a/src/services/llm-utils/summarizer.ts b/src/services/llm-utils/summarizer.ts new file mode 100644 index 0000000000..8298d3d29e --- /dev/null +++ b/src/services/llm-utils/summarizer.ts @@ -0,0 +1,209 @@ +import { z } from "zod" +import { LlmClient } from "./llm-client" +import { SummaryResult } from "./types" +import { CodeBlock } from "../code-index/interfaces" + +// Schema for summarization response +const summaryResponseSchema = z.object({ + title: z.string().optional(), + summary: z.string(), + tags: z.array(z.string()).optional(), +}) + +/** + * Generates summaries and tags for code blocks + */ +export class Summarizer { + constructor(private readonly llmClient: LlmClient) {} + + /** + * Generate a summary for a code block + */ + async summarize(codeBlock: CodeBlock, options: SummarizeOptions = {}): Promise { + const systemPrompt = `You are a code documentation expert. Generate concise summaries for code blocks. + +Requirements: +- Title: One-line description (optional, only for functions/classes) +- Summary: 1-2 sentences describing what the code does +- Tags: 2-5 relevant keywords for searchability + +Focus on: +- Main functionality and purpose +- Key algorithms or patterns used +- Dependencies and relationships +- Error handling or edge cases + +Be concise and factual. Respond with JSON only.` + + const userPrompt = this.buildPrompt(codeBlock, options) + + try { + const response = await this.llmClient.generateJson>(userPrompt, { + schema: summaryResponseSchema, + systemPrompt, + maxTokens: 300, + temperature: 0.2, + }) + + return response + } catch (error) { + console.error("[Summarizer] Failed to generate summary:", error) + // Fallback to empty summary + return { + summary: "Code block", + tags: [], + } + } + } + + /** + * Build the prompt for summarization + */ + private buildPrompt(codeBlock: CodeBlock, options: SummarizeOptions): string { + const language = options.language || this.detectLanguage(codeBlock.file_path) + const codeSnippet = this.truncateCode(codeBlock.content, options.maxCodeLength || 500) + + return `Summarize this ${language} code: + +File: ${codeBlock.file_path} +Lines: ${codeBlock.start_line}-${codeBlock.end_line} +${codeBlock.identifier ? `Identifier: ${codeBlock.identifier}` : ""} + +Code: +\`\`\`${language} +${codeSnippet} +\`\`\` + +Generate a JSON response with: +- title (optional): One-line description if it's a function/class +- summary: 1-2 sentences about what this code does +- tags: 2-5 relevant keywords + +Example: +{ + "title": "Calculate user authentication token", + "summary": "Generates a JWT token for user authentication with expiration and refresh logic.", + "tags": ["auth", "jwt", "token", "security"] +}` + } + + /** + * Batch summarize multiple code blocks + */ + async summarizeBatch(codeBlocks: CodeBlock[], options: SummarizeOptions = {}): Promise> { + const results = new Map() + + // Process in parallel with concurrency limit + const batchSize = options.batchSize || 5 + for (let i = 0; i < codeBlocks.length; i += batchSize) { + const batch = codeBlocks.slice(i, i + batchSize) + const batchResults = await Promise.all(batch.map((block) => this.summarize(block, options))) + + batch.forEach((block, index) => { + // Use segmentHash as unique identifier + results.set(block.segmentHash, batchResults[index]) + }) + } + + return results + } + + /** + * Detect language from file path + */ + private detectLanguage(filePath: string): string { + const ext = filePath.split(".").pop()?.toLowerCase() + + const languageMap: Record = { + ts: "typescript", + tsx: "typescript", + js: "javascript", + jsx: "javascript", + py: "python", + java: "java", + cpp: "cpp", + c: "c", + cs: "csharp", + go: "go", + rs: "rust", + rb: "ruby", + php: "php", + swift: "swift", + kt: "kotlin", + scala: "scala", + r: "r", + sql: "sql", + sh: "bash", + yaml: "yaml", + yml: "yaml", + json: "json", + xml: "xml", + html: "html", + css: "css", + scss: "scss", + sass: "sass", + less: "less", + } + + return languageMap[ext || ""] || "text" + } + + /** + * Truncate code to maximum length + */ + private truncateCode(code: string, maxLength: number): string { + if (code.length <= maxLength) { + return code + } + + // Try to truncate at a natural boundary + const truncated = code.substring(0, maxLength) + const lastNewline = truncated.lastIndexOf("\n") + + if (lastNewline > maxLength * 0.8) { + return truncated.substring(0, lastNewline) + "\n// ..." + } + + return truncated + "..." + } + + /** + * Generate augmented embedding text with summary and tags + */ + generateAugmentedText(codeBlock: CodeBlock, summary: SummaryResult): string { + const language = this.detectLanguage(codeBlock.file_path) + const relPath = codeBlock.file_path + const identifier = codeBlock.identifier || "" + const span = `${codeBlock.start_line}-${codeBlock.end_line}` + + // Build augmented text for embedding + const parts = [ + language, + relPath, + identifier, + span, + summary.title || "", + summary.summary, + ...(summary.tags || []), + codeBlock.content, + ].filter(Boolean) + + return parts.join(" | ") + } +} + +/** + * Options for summarization + */ +export interface SummarizeOptions { + /** Programming language override */ + language?: string + /** Maximum code length to process */ + maxCodeLength?: number + /** Batch size for parallel processing */ + batchSize?: number + /** Focus on specific aspects */ + focus?: "functionality" | "api" | "implementation" + /** Include specific metadata */ + includeMetadata?: boolean +} diff --git a/src/services/llm-utils/types.ts b/src/services/llm-utils/types.ts new file mode 100644 index 0000000000..f06f1f356e --- /dev/null +++ b/src/services/llm-utils/types.ts @@ -0,0 +1,83 @@ +import { ProviderSettings } from "@roo-code/types" + +/** + * Configuration for the sub-LLM system + */ +export interface SubLlmConfig { + /** Whether sub-LLM features are enabled */ + enabled: boolean + /** Mode for model selection: 'mirror' uses chat model, 'custom' allows override */ + modelMode: SubLlmMode + /** Provider settings when using custom mode */ + customProvider?: ProviderSettings + /** Maximum tokens per operation */ + maxTokensPerOp?: number + /** Daily cost cap in USD */ + dailyCostCapUSD?: number + /** Timeout for LLM operations in milliseconds */ + timeout?: number +} + +export type SubLlmMode = "mirror" | "custom" + +/** + * Query variant for multi-query expansion + */ +export interface QueryVariant { + /** The rewritten query */ + query: string + /** Type of variant (e.g., 'synonym', 'symbol', 'natural') */ + type: string + /** Optional explanation of the rewrite */ + reason?: string +} + +/** + * Result from reranking operation + */ +export interface RerankResult { + /** Original item ID or index */ + id: string | number + /** Rerank score between 0 and 1 */ + score: number + /** Optional reasoning for the score */ + reason?: string +} + +/** + * Result from summarization operation + */ +export interface SummaryResult { + /** One-line title */ + title?: string + /** 1-2 sentence summary */ + summary: string + /** 2-5 relevant tags */ + tags?: string[] +} + +/** + * Options for LLM generation + */ +export interface GenerateOptions { + /** Maximum tokens for the response */ + maxTokens?: number + /** Temperature for generation */ + temperature?: number + /** System prompt override */ + systemPrompt?: string + /** Timeout in milliseconds */ + timeout?: number +} + +/** + * JSON generation options with schema validation + */ +export interface GenerateJsonOptions extends GenerateOptions { + /** Zod schema for validation */ + schema?: any + /** Whether to retry on validation failure */ + retryOnValidationFailure?: boolean + /** Maximum retry attempts */ + maxRetries?: number +}