entity-extractor.ts•3.31 kB
import { AIContext } from '../types/index.js';
interface ExtractedEntities {
primary_target: string;
parameters: Record<string, any>;
confidence: number;
}
export class EntityExtractor {
private deviceNameMap: Map<string, string>;
private parameterPatterns: Map<string, RegExp>;
constructor() {
this.deviceNameMap = new Map();
this.parameterPatterns = new Map();
this.initializePatterns();
}
private initializePatterns(): void {
// Device name variations
this.deviceNameMap.set('living room light', 'light.living_room');
this.deviceNameMap.set('kitchen light', 'light.kitchen');
this.deviceNameMap.set('bedroom light', 'light.bedroom');
// Parameter patterns
this.parameterPatterns.set('brightness', /(\d+)\s*(%|percent)|bright(ness)?\s+(\d+)/i);
this.parameterPatterns.set('temperature', /(\d+)\s*(degrees?|°)[CF]?/i);
this.parameterPatterns.set('color', /(red|green|blue|white|warm|cool)/i);
}
async extract(input: string): Promise<ExtractedEntities> {
const entities: ExtractedEntities = {
primary_target: '',
parameters: {},
confidence: 0
};
try {
// Find device name
for (const [key, value] of this.deviceNameMap) {
if (input.toLowerCase().includes(key)) {
entities.primary_target = value;
break;
}
}
// Extract parameters
for (const [param, pattern] of this.parameterPatterns) {
const match = input.match(pattern);
if (match) {
entities.parameters[param] = this.normalizeParameterValue(param, match[1]);
}
}
// Calculate confidence based on matches
entities.confidence = this.calculateConfidence(entities, input);
return entities;
} catch (error) {
console.error('Entity extraction error:', error);
return {
primary_target: '',
parameters: {},
confidence: 0
};
}
}
private normalizeParameterValue(parameter: string, value: string): number | string {
switch (parameter) {
case 'brightness':
return Math.min(100, Math.max(0, parseInt(value)));
case 'temperature':
return parseInt(value);
case 'color':
return value.toLowerCase();
default:
return value;
}
}
private calculateConfidence(entities: ExtractedEntities, input: string): number {
let confidence = 0;
// Device confidence
if (entities.primary_target) {
confidence += 0.5;
}
// Parameter confidence
const paramCount = Object.keys(entities.parameters).length;
confidence += paramCount * 0.25;
// Normalize confidence to 0-1 range
return Math.min(1, confidence);
}
async updateDeviceMap(devices: Record<string, string>): Promise<void> {
for (const [key, value] of Object.entries(devices)) {
this.deviceNameMap.set(key, value);
}
}
}