99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from app.settings import Settings
|
|
|
|
|
|
TOKEN_PATTERN = re.compile(r"[a-zA-Z0-9_]+")
|
|
CN_TOKEN_PATTERN = re.compile(r"[\u4e00-\u9fff]")
|
|
|
|
TERM_ALIASES = {
|
|
"法兰": {"flange", "mounting_plate"},
|
|
"套筒": {"sleeve", "tube", "cylinder"},
|
|
"孔": {"hole", "circle", "circles"},
|
|
"沉头": {"countersunk", "counterbore"},
|
|
"沉孔": {"counterbore", "counterbored"},
|
|
"槽": {"slot", "obround"},
|
|
"底板": {"plate", "rectangle", "mounting_plate"},
|
|
"轴": {"shaft", "cylinder"},
|
|
}
|
|
|
|
|
|
def tokens(value: str) -> set[str]:
|
|
text = str(value or "").lower()
|
|
result = {token.lower() for token in TOKEN_PATTERN.findall(text)}
|
|
result.update(CN_TOKEN_PATTERN.findall(text))
|
|
for term, aliases in TERM_ALIASES.items():
|
|
if term in text:
|
|
result.update(aliases)
|
|
return result
|
|
|
|
|
|
class CdslLibrary:
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
self.index_path = settings.library_root / "index" / "catalog.json"
|
|
|
|
def _records(self) -> list[dict[str, Any]]:
|
|
if not self.index_path.is_file():
|
|
return []
|
|
payload = json.loads(self.index_path.read_text(encoding="utf-8"))
|
|
return payload.get("samples", [])
|
|
|
|
def count(self) -> int:
|
|
return len(self._records())
|
|
|
|
def search(self, query: str, limit: int = 5) -> list[dict[str, Any]]:
|
|
query_tokens = tokens(query)
|
|
if not query_tokens:
|
|
return [
|
|
{
|
|
"part_id": record["part_id"],
|
|
"profiles": record.get("profiles", []),
|
|
"features": record.get("features", []),
|
|
"summary": record.get("summary", ""),
|
|
"family": record.get("family", ""),
|
|
"atomics": record.get("features", []),
|
|
}
|
|
for record in self._records()[:limit]
|
|
]
|
|
scored: list[tuple[int, dict[str, Any]]] = []
|
|
for record in self._records():
|
|
corpus = " ".join([
|
|
record.get("part_id", ""),
|
|
record.get("source_name", ""),
|
|
record.get("summary", ""),
|
|
record.get("family", ""),
|
|
" ".join(record.get("profiles", [])),
|
|
" ".join(record.get("features", [])),
|
|
" ".join(record.get("parameters", [])),
|
|
])
|
|
score = len(query_tokens & tokens(corpus))
|
|
if score:
|
|
scored.append((score, record))
|
|
scored.sort(key=lambda item: (-item[0], item[1]["part_id"]))
|
|
return [
|
|
{
|
|
"part_id": record["part_id"],
|
|
"profiles": record.get("profiles", []),
|
|
"features": record.get("features", []),
|
|
"summary": record.get("summary", ""),
|
|
"family": record.get("family", ""),
|
|
"atomics": record.get("features", []),
|
|
}
|
|
for _, record in scored[:limit]
|
|
]
|
|
|
|
def read_sample(self, part_id: str) -> dict[str, Any]:
|
|
for record in self._records():
|
|
if record.get("part_id") == part_id:
|
|
source = self.settings.library_root / "samples" / part_id / "model.cdsl.json"
|
|
if not source.is_file():
|
|
break
|
|
return json.loads(source.read_text(encoding="utf-8"))
|
|
raise ValueError(f"CDSL sample not found: {part_id}")
|