mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-06 08:18:39 +00:00
feat: implement LLM-assisted code index search with shared sub-LLM layer
- Add shared sub-LLM layer in src/services/llm-utils/ - LlmClient for provider selection and API methods with budget management - JsonRunner for JSON extraction and validation - QueryRewriter for multi-query expansion - Reranker for LLM-based result scoring - Summarizer for code block summaries and tags - Integrate query rewriting and reranking in search - Create EnhancedCodeIndexSearchService with optional LLM features - Support multi-query search with result merging - Blend reranked scores with embedding scores - Add optional summary/tag generation - Create EnhancedScannerFactory for conditional enhancement - Support augmented embeddings with summaries - Add comprehensive configuration settings - Sub-LLM enable/disable and model selection - Individual feature toggles for rewriter, reranker, summaries, tags - Budget controls and timeout settings All features are opt-in and disabled by default to ensure no additional costs unless explicitly enabled by users.
This commit is contained in:
Roo Code
parent
247da38b02
commit
a82678aec6
10 changed files with 1557 additions and 0 deletions
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
246
src/services/code-index/enhanced-search-service.ts
Normal file
246
src/services/code-index/enhanced-search-service.ts
Normal file
|
|
@ -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<void> {
|
||||
try {
|
||||
// Get sub-LLM configuration from VSCode settings
|
||||
const config = vscode.workspace.getConfiguration(Package.name)
|
||||
|
||||
this.subLlmConfig = {
|
||||
enabled: config.get<boolean>("subLlm.enabled", false),
|
||||
modelMode: config.get<"mirror" | "custom">("subLlm.model.mode", "mirror"),
|
||||
maxTokensPerOp: config.get<number>("subLlm.maxTokensPerOp", 1000),
|
||||
dailyCostCapUSD: config.get<number>("subLlm.dailyCostCapUSD", 1.0),
|
||||
timeout: config.get<number>("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<VectorStoreSearchResult[]> {
|
||||
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<string[]> {
|
||||
const queries = [query] // Always include original
|
||||
|
||||
// Check if query rewriting is enabled
|
||||
const rewriterEnabled = vscode.workspace
|
||||
.getConfiguration(Package.name)
|
||||
.get<boolean>("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<VectorStoreSearchResult[][]> {
|
||||
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<string | number, VectorStoreSearchResult>()
|
||||
|
||||
// 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<VectorStoreSearchResult[]> {
|
||||
// Check if reranking is enabled
|
||||
const rerankerEnabled = vscode.workspace
|
||||
.getConfiguration(Package.name)
|
||||
.get<boolean>("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<number>("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,
|
||||
}
|
||||
}
|
||||
}
|
||||
185
src/services/code-index/processors/enhanced-scanner.ts
Normal file
185
src/services/code-index/processors/enhanced-scanner.ts
Normal file
|
|
@ -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<IDirectoryScanner> {
|
||||
const config = vscode.workspace.getConfiguration(Package.name)
|
||||
|
||||
// Check if LLM features are enabled
|
||||
const summariesEnabled = config.get<boolean>("codeIndex.llm.summaries", false)
|
||||
const tagsEnabled = config.get<boolean>("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<void> {
|
||||
try {
|
||||
const config = vscode.workspace.getConfiguration(Package.name)
|
||||
|
||||
// Initialize sub-LLM configuration
|
||||
const subLlmConfig: SubLlmConfig = {
|
||||
enabled: config.get<boolean>("subLlm.enabled", false),
|
||||
modelMode: config.get<"mirror" | "custom">("subLlm.model.mode", "mirror"),
|
||||
maxTokensPerOp: config.get<number>("subLlm.maxTokensPerOp", 500),
|
||||
dailyCostCapUSD: config.get<number>("subLlm.dailyCostCapUSD", 1.0),
|
||||
timeout: config.get<number>("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()
|
||||
}
|
||||
}
|
||||
6
src/services/llm-utils/index.ts
Normal file
6
src/services/llm-utils/index.ts
Normal file
|
|
@ -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"
|
||||
199
src/services/llm-utils/json-runner.ts
Normal file
199
src/services/llm-utils/json-runner.ts
Normal file
|
|
@ -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<T>(text: string, schema?: z.ZodSchema<T>): 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<T>(
|
||||
data: unknown,
|
||||
schema: z.ZodSchema<T>,
|
||||
): { 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 }
|
||||
}
|
||||
}
|
||||
254
src/services/llm-utils/llm-client.ts
Normal file
254
src/services/llm-utils/llm-client.ts
Normal file
|
|
@ -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<void> {
|
||||
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<string> {
|
||||
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<T>(prompt: string, options: GenerateJsonOptions<T> = {}): Promise<T> {
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
113
src/services/llm-utils/query-rewriter.ts
Normal file
113
src/services/llm-utils/query-rewriter.ts
Normal file
|
|
@ -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<QueryVariant[]> {
|
||||
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<z.infer<typeof queryRewriteSchema>>(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<Map<string, QueryVariant[]>> {
|
||||
const results = new Map<string, QueryVariant[]>()
|
||||
|
||||
// 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
|
||||
}
|
||||
190
src/services/llm-utils/reranker.ts
Normal file
190
src/services/llm-utils/reranker.ts
Normal file
|
|
@ -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<RerankResult[]> {
|
||||
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<z.infer<typeof rerankResponseSchema>>(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<string | number, RerankResult>()
|
||||
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
|
||||
}
|
||||
}
|
||||
209
src/services/llm-utils/summarizer.ts
Normal file
209
src/services/llm-utils/summarizer.ts
Normal file
|
|
@ -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<SummaryResult> {
|
||||
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<z.infer<typeof summaryResponseSchema>>(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<Map<string, SummaryResult>> {
|
||||
const results = new Map<string, SummaryResult>()
|
||||
|
||||
// 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<string, string> = {
|
||||
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
|
||||
}
|
||||
83
src/services/llm-utils/types.ts
Normal file
83
src/services/llm-utils/types.ts
Normal file
|
|
@ -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<T> extends GenerateOptions {
|
||||
/** Zod schema for validation */
|
||||
schema?: any
|
||||
/** Whether to retry on validation failure */
|
||||
retryOnValidationFailure?: boolean
|
||||
/** Maximum retry attempts */
|
||||
maxRetries?: number
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue