zrald
Version:
Advanced Graph RAG MCP Server with sophisticated graph structures, operators, and agentic capabilities for AI agents
411 lines • 15.4 kB
JavaScript
import { v4 as uuidv4 } from 'uuid';
export class QueryPlanner {
graphDb;
vectorStore;
planCache = new Map();
constructor(graphDb, vectorStore) {
this.graphDb = graphDb;
this.vectorStore = vectorStore;
}
async createQueryPlan(query, context) {
// Check cache first
const cacheKey = this.generateCacheKey(query, context);
const cachedPlan = this.planCache.get(cacheKey);
if (cachedPlan) {
console.log('Using cached query plan');
return cachedPlan;
}
// Analyze the query
const analysis = await this.analyzeQuery(query, context);
// Create operator chain based on analysis
const operatorChain = await this.createOperatorChain(analysis, query, context);
// Estimate execution cost
const estimatedCost = this.estimateExecutionCost(operatorChain, analysis);
// Determine priority
const priority = this.determinePriority(analysis, context);
const plan = {
id: uuidv4(),
query,
intent: analysis.intent,
operator_chain: operatorChain,
estimated_cost: estimatedCost,
priority
};
// Cache the plan
this.planCache.set(cacheKey, plan);
return plan;
}
async analyzeQuery(query, context) {
const queryLower = query.toLowerCase();
// Determine intent
let intent = 'factual';
if (queryLower.includes('compare') || queryLower.includes('difference') || queryLower.includes('versus')) {
intent = 'comparative';
}
else if (queryLower.includes('analyze') || queryLower.includes('relationship') || queryLower.includes('connection')) {
intent = 'analytical';
}
else if (queryLower.includes('explore') || queryLower.includes('discover') || queryLower.includes('find all')) {
intent = 'exploratory';
}
// Extract entities and concepts (simplified NER)
const entities = await this.extractEntities(query);
const concepts = await this.extractConcepts(query);
// Determine complexity
const complexity = this.determineComplexity(query, entities, concepts);
// Estimate result size
const expectedResultSize = this.estimateResultSize(intent, complexity, entities.length);
// Check if reasoning is required
const requiresReasoning = this.requiresReasoning(query, intent);
// Check for temporal aspects
const temporalAspects = this.hasTemporalAspects(query);
return {
intent,
entities,
concepts,
complexity,
expectedResultSize,
requiresReasoning,
temporalAspects
};
}
async extractEntities(query) {
// Simplified entity extraction - in practice, use NER
const entities = [];
// Look for quoted strings as potential entities
const quotedMatches = query.match(/"([^"]+)"/g);
if (quotedMatches) {
entities.push(...quotedMatches.map(match => match.slice(1, -1)));
}
// Look for capitalized words as potential entities
const capitalizedWords = query.match(/\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b/g);
if (capitalizedWords) {
entities.push(...capitalizedWords);
}
return [...new Set(entities)];
}
async extractConcepts(query) {
// Simplified concept extraction
const conceptKeywords = [
'technology', 'science', 'business', 'health', 'education',
'environment', 'politics', 'economics', 'culture', 'society'
];
const concepts = [];
const queryLower = query.toLowerCase();
for (const concept of conceptKeywords) {
if (queryLower.includes(concept)) {
concepts.push(concept);
}
}
return concepts;
}
determineComplexity(query, entities, concepts) {
let complexityScore = 0;
// Length factor
if (query.length > 100)
complexityScore += 1;
if (query.length > 200)
complexityScore += 1;
// Entity count factor
if (entities.length > 2)
complexityScore += 1;
if (entities.length > 5)
complexityScore += 1;
// Concept count factor
if (concepts.length > 1)
complexityScore += 1;
// Complexity keywords
const complexKeywords = ['relationship', 'connection', 'pattern', 'trend', 'correlation', 'causation'];
for (const keyword of complexKeywords) {
if (query.toLowerCase().includes(keyword)) {
complexityScore += 1;
break;
}
}
if (complexityScore <= 1)
return 'low';
if (complexityScore <= 3)
return 'medium';
return 'high';
}
estimateResultSize(intent, complexity, entityCount) {
if (intent === 'factual' && complexity === 'low')
return 'small';
if (intent === 'exploratory' || complexity === 'high')
return 'large';
if (entityCount > 3)
return 'large';
return 'medium';
}
requiresReasoning(query, intent) {
const reasoningKeywords = ['why', 'how', 'explain', 'reason', 'cause', 'effect', 'implication'];
const queryLower = query.toLowerCase();
return intent === 'analytical' ||
intent === 'comparative' ||
reasoningKeywords.some(keyword => queryLower.includes(keyword));
}
hasTemporalAspects(query) {
const temporalKeywords = ['when', 'before', 'after', 'during', 'timeline', 'history', 'evolution', 'trend'];
const queryLower = query.toLowerCase();
return temporalKeywords.some(keyword => queryLower.includes(keyword));
}
async createOperatorChain(analysis, query, context) {
const operators = [];
// Always start with vector search for semantic relevance
operators.push({
type: 'VDBOperator',
config: {
query_embedding: await this.getQueryEmbedding(query),
top_k: this.getTopKForAnalysis(analysis),
similarity_threshold: this.getThresholdForAnalysis(analysis),
node_types: this.getRelevantNodeTypes(analysis)
},
dependencies: []
});
// Add relationship exploration based on intent
if (analysis.intent === 'analytical' || analysis.intent === 'exploratory') {
operators.push({
type: 'OneHopOperator',
config: {
source_nodes: [], // Will be populated from VDB results
direction: 'both',
max_depth: analysis.complexity === 'high' ? 2 : 1
},
dependencies: ['VDBOperator']
});
}
// Add PageRank for authority analysis in complex queries
if (analysis.complexity === 'high' || analysis.requiresReasoning) {
operators.push({
type: 'PPROperator',
config: {
seed_nodes: [], // Will be populated from previous results
damping_factor: 0.85,
max_iterations: 100
},
dependencies: ['VDBOperator']
});
}
// Add path finding for comparative or reasoning queries
if (analysis.intent === 'comparative' || analysis.requiresReasoning) {
operators.push({
type: 'KHopPathOperator',
config: {
source_nodes: [],
target_nodes: [],
max_hops: 3,
path_limit: 10
},
dependencies: ['VDBOperator', 'OneHopOperator']
});
}
// Add co-occurrence analysis for entity-rich queries
if (analysis.entities.length >= 2) {
operators.push({
type: 'OccurrenceOperator',
config: {
entities: [], // Will be populated with entity IDs
co_occurrence_window: 50,
min_frequency: 1
},
dependencies: ['VDBOperator']
});
}
// Add aggregation for complex multi-faceted queries
if (analysis.complexity === 'high' && operators.length > 2) {
operators.push({
type: 'AggregatorOperator',
config: {
source_nodes: [],
aggregation_method: 'weighted',
normalize_scores: true
},
dependencies: operators.slice(0, -1).map(op => op.type)
});
}
// Determine execution pattern
const executionPattern = this.determineExecutionPattern(analysis, operators.length);
// Determine fusion strategy
const fusionStrategy = this.determineFusionStrategy(analysis);
return {
operators,
execution_pattern: executionPattern,
fusion_strategy: fusionStrategy
};
}
async getQueryEmbedding(query) {
// In practice, this would use a proper embedding model
// For now, return a dummy embedding
return Array(384).fill(0).map(() => Math.random());
}
getTopKForAnalysis(analysis) {
switch (analysis.expectedResultSize) {
case 'small': return 5;
case 'medium': return 15;
case 'large': return 30;
default: return 10;
}
}
getThresholdForAnalysis(analysis) {
switch (analysis.complexity) {
case 'low': return 0.8;
case 'medium': return 0.7;
case 'high': return 0.6;
default: return 0.7;
}
}
getRelevantNodeTypes(analysis) {
const nodeTypes = [];
if (analysis.entities.length > 0) {
nodeTypes.push('entity');
}
if (analysis.concepts.length > 0) {
nodeTypes.push('concept');
}
if (analysis.intent === 'factual') {
nodeTypes.push('document', 'chunk');
}
return nodeTypes.length > 0 ? nodeTypes : ['entity', 'concept', 'document'];
}
determineExecutionPattern(analysis, operatorCount) {
if (analysis.complexity === 'high' && operatorCount > 3) {
return 'adaptive';
}
if (operatorCount <= 2 || analysis.intent === 'factual') {
return 'sequential';
}
return 'parallel';
}
determineFusionStrategy(analysis) {
if (analysis.intent === 'comparative') {
return 'intersection';
}
if (analysis.complexity === 'high' || analysis.requiresReasoning) {
return 'weighted_average';
}
return 'union';
}
estimateExecutionCost(operatorChain, analysis) {
let cost = 0;
// Base cost per operator
cost += operatorChain.operators.length * 10;
// Complexity multiplier
const complexityMultiplier = {
'low': 1,
'medium': 2,
'high': 4
};
cost *= complexityMultiplier[analysis.complexity];
// Execution pattern cost
const patternMultiplier = {
'sequential': 1,
'parallel': 1.5,
'adaptive': 2
};
cost *= patternMultiplier[operatorChain.execution_pattern];
// Expected result size impact
const sizeMultiplier = {
'small': 1,
'medium': 1.5,
'large': 2.5
};
cost *= sizeMultiplier[analysis.expectedResultSize];
return Math.round(cost);
}
determinePriority(analysis, context) {
// Priority based on context and analysis
if (context?.urgent)
return 'critical';
if (analysis.complexity === 'high')
return 'high';
if (analysis.intent === 'factual' && analysis.complexity === 'low')
return 'low';
return 'medium';
}
generateCacheKey(query, context) {
const contextStr = context ? JSON.stringify(context) : '';
return Buffer.from(query + contextStr).toString('base64');
}
// Plan optimization methods
async optimizePlan(plan) {
const optimizedChain = await this.optimizeOperatorChain(plan.operator_chain);
return {
...plan,
operator_chain: optimizedChain,
estimated_cost: this.estimateExecutionCost(optimizedChain, {
complexity: 'medium',
expectedResultSize: 'medium'
})
};
}
async optimizeOperatorChain(chain) {
// Remove redundant operators
const optimizedOperators = this.removeRedundantOperators(chain.operators);
// Reorder for better performance
const reorderedOperators = this.reorderOperators(optimizedOperators);
// Optimize configurations
const optimizedConfigs = await this.optimizeOperatorConfigs(reorderedOperators);
return {
...chain,
operators: optimizedConfigs
};
}
removeRedundantOperators(operators) {
// Remove duplicate operator types with similar configs
const seen = new Map();
const filtered = [];
for (const op of operators) {
const key = `${op.type}_${JSON.stringify(op.config)}`;
if (!seen.has(key)) {
seen.set(key, op);
filtered.push(op);
}
}
return filtered;
}
reorderOperators(operators) {
// Reorder operators for optimal execution
// Put cheaper operations first, expensive ones later
const operatorCosts = {
'VDBOperator': 1,
'OneHopOperator': 2,
'PPROperator': 4,
'KHopPathOperator': 3,
'OccurrenceOperator': 2,
'AggregatorOperator': 1
};
return operators.sort((a, b) => {
const costA = operatorCosts[a.type] || 5;
const costB = operatorCosts[b.type] || 5;
return costA - costB;
});
}
async optimizeOperatorConfigs(operators) {
// Optimize individual operator configurations
return operators.map(op => {
const optimizedConfig = { ...op.config };
// Optimize top_k values
if (optimizedConfig.top_k && optimizedConfig.top_k > 50) {
optimizedConfig.top_k = 50; // Cap at reasonable limit
}
// Optimize thresholds
if (optimizedConfig.similarity_threshold && optimizedConfig.similarity_threshold < 0.3) {
optimizedConfig.similarity_threshold = 0.3; // Prevent too low thresholds
}
return {
...op,
config: optimizedConfig
};
});
}
clearCache() {
this.planCache.clear();
}
getCacheStats() {
return {
size: this.planCache.size,
keys: Array.from(this.planCache.keys())
};
}
}
//# sourceMappingURL=query-planner.js.map