Files
tfm_ainventory/backend/services/spec_extractor.py

223 lines
7.5 KiB
Python

"""
Specification extractor service for search results parsing.
Extracts product specifications (manufacturer, model, capacity, specs) from search
results and maps them to Item model fields for pre-population in onboarding UI.
"""
import re
from dataclasses import dataclass, asdict
from typing import Optional, Dict, Any
import logging
log = logging.getLogger("ainventory")
@dataclass
class ExtractedSpecs:
"""Extracted specifications from search results."""
manufacturer: Optional[str] = None
model: Optional[str] = None
capacity: Optional[str] = None # e.g., "16GB"
memory_type: Optional[str] = None # e.g., "DDR4"
speed: Optional[str] = None # e.g., "3200MHz"
latency: Optional[str] = None # e.g., "CAS 16"
storage_type: Optional[str] = None # e.g., "SSD", "HDD"
processor_brand: Optional[str] = None # e.g., "Intel"
processor_model: Optional[str] = None # e.g., "Core i7-12700K"
power_rating: Optional[str] = None # e.g., "850W"
description: str = "" # Full snippet/details from search
confidence: float = 0.0 # 0.0-1.0 score
def to_item_fields(self, category: str) -> Dict[str, str]:
"""
Map extracted specs to Item model fields.
Args:
category: Item category (e.g., "Memory", "Storage", "Processor")
Returns:
Dict with keys: type, description, notes
"""
type_str = ""
notes_parts = []
if category.lower() in ["memory", "ram"]:
if self.memory_type:
type_str = self.memory_type
if self.capacity:
notes_parts.append(f"Capacity: {self.capacity}")
if self.speed:
notes_parts.append(f"Speed: {self.speed}")
if self.latency:
notes_parts.append(f"Latency: {self.latency}")
elif category.lower() in ["storage", "ssd", "hdd"]:
if self.storage_type:
type_str = self.storage_type
if self.capacity:
notes_parts.append(f"Capacity: {self.capacity}")
if self.model:
notes_parts.append(f"Model: {self.model}")
elif category.lower() in ["processor", "cpu", "gpu"]:
if self.processor_brand:
type_str = self.processor_brand
if self.processor_model:
notes_parts.append(f"Model: {self.processor_model}")
elif category.lower() in ["power", "psu"]:
if self.power_rating:
type_str = self.power_rating
if self.model:
notes_parts.append(f"Model: {self.model}")
if self.manufacturer and self.manufacturer not in type_str:
notes_parts.insert(0, f"Manufacturer: {self.manufacturer}")
return {
"type": type_str or "Generic",
"description": self.description[:200] if self.description else "",
"notes": " | ".join(notes_parts) if notes_parts else ""
}
def extract_specs_from_search(title: str, snippet: str, url: str = "") -> ExtractedSpecs:
"""
Extract specifications from a search result (title + snippet).
Args:
title: Search result title
snippet: Search result snippet/description
url: Source URL (optional)
Returns:
ExtractedSpecs object with extracted fields and confidence score
"""
full_text = f"{title} {snippet}".upper()
specs = ExtractedSpecs(description=snippet or title)
confidence_score = 0.0
# Manufacturer extraction
manufacturers = ["KINGSTON", "SAMSUNG", "CORSAIR", "INTEL", "AMD", "NVIDIA",
"CRUCIAL", "ADATA", "WESTERN DIGITAL", "SEAGATE", "HP", "DELL",
"LENOVO", "ASUS", "GIGABYTE", "MSI", "EVGA", "SAPPHIRE"]
for mfg in manufacturers:
if mfg in full_text:
specs.manufacturer = mfg.title()
confidence_score += 15
break
# Memory type extraction (DDR3/4/5, SODIMM, DIMM)
memory_patterns = {
r"DDR5?(\s|-)?LP?": "DDR5",
r"DDR4\b": "DDR4",
r"DDR3\b": "DDR3",
r"SODIMM\b": "SODIMM",
r"DIMM\b": "DIMM",
}
for pattern, memory_type in memory_patterns.items():
if re.search(pattern, full_text):
specs.memory_type = memory_type
confidence_score += 20
break
# Capacity extraction (16GB, 512GB, 1TB, etc.)
capacity_match = re.search(r"(\d+(?:\.\d+)?)\s*(GB|TB|MB)", full_text)
if capacity_match:
specs.capacity = f"{capacity_match.group(1)}{capacity_match.group(2)}"
confidence_score += 25
# Speed extraction (3200MHz, 6400MT/s, etc.)
speed_match = re.search(r"(\d+)\s*(MHz|MT/S)", full_text)
if speed_match:
specs.speed = f"{speed_match.group(1)}{speed_match.group(2)}"
confidence_score += 10
# Latency extraction (CAS 16, CAS 18, etc.)
latency_match = re.search(r"CAS\s*(\d+)", full_text)
if latency_match:
specs.latency = f"CAS {latency_match.group(1)}"
confidence_score += 10
# Storage type extraction (SSD, HDD, NVMe, M.2)
storage_patterns = {
r"NVME\b|NVMe\b": "NVMe",
r"SSD\b": "SSD",
r"HDD\b|HARD DRIVE": "HDD",
r"M\.2\b": "M.2",
}
for pattern, storage_type in storage_patterns.items():
if re.search(pattern, full_text):
specs.storage_type = storage_type
confidence_score += 20
break
# Processor extraction
processor_patterns = {
r"INTEL\s+(CORE\s+)?(I[3579]|PENTIUM|CELERON|XEON)": "Intel",
r"AMD\s+(RYZEN|EPYC|FX)": "AMD",
r"NVIDIA\s+GEFORCE": "NVIDIA",
}
for pattern, brand in processor_patterns.items():
if re.search(pattern, full_text):
specs.processor_brand = brand
confidence_score += 15
break
# Model number extraction (alphanumeric patterns after brand)
if specs.manufacturer:
# Look for patterns like "Kingston KF466C40RS-16" or "Samsung 870 EVO"
model_match = re.search(rf"{specs.manufacturer.upper()}\s+([A-Z0-9\-]+)", full_text)
if model_match:
specs.model = model_match.group(1)
confidence_score += 20
# Power rating extraction (850W, 1000W, etc.)
power_match = re.search(r"(\d+)(\s*)W(?=\s|$|\D)", full_text)
if power_match:
specs.power_rating = f"{power_match.group(1)}W"
confidence_score += 15
# Normalize confidence to 0.0-1.0 range
specs.confidence = min(100.0, confidence_score) / 100.0
return specs
def extract_specs_from_multiple_results(
results: list[Dict[str, str]],
category: str
) -> Dict[str, str]:
"""
Extract specs from multiple search results and return best candidate fields.
Aggregates specifications across multiple results, preferring specs with
highest confidence and deduplication.
Args:
results: List of search result dicts with 'title', 'snippet', 'url'
category: Item category for field mapping
Returns:
Dict with keys: type, description, notes
"""
if not results:
return {"type": "", "description": "", "notes": ""}
# Extract from all results and pick highest confidence
all_specs = []
for result in results:
specs = extract_specs_from_search(
result.get("title", ""),
result.get("snippet", ""),
result.get("url", "")
)
all_specs.append(specs)
# Pick spec set with highest confidence
best_specs = max(all_specs, key=lambda s: s.confidence) if all_specs else ExtractedSpecs()
return best_specs.to_item_fields(category)