|
|
from __future__ import annotations |
|
|
|
|
|
import json |
|
|
from dataclasses import dataclass |
|
|
from datetime import datetime |
|
|
from pathlib import Path |
|
|
from typing import List, Tuple |
|
|
|
|
|
import faiss |
|
|
import numpy as np |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class VectorIndexConfig: |
|
|
index_type: str = "IndexFlatIP" |
|
|
embedding_dim: int = 384 |
|
|
|
|
|
|
|
|
class VectorIndex: |
|
|
def __init__(self, embeddings: np.ndarray, config: VectorIndexConfig | None = None): |
|
|
cfg = config or VectorIndexConfig(embedding_dim=embeddings.shape[1]) |
|
|
if cfg.index_type == "IndexFlatIP": |
|
|
self.index = faiss.IndexFlatIP(cfg.embedding_dim) |
|
|
elif cfg.index_type == "IndexHNSWFlat": |
|
|
self.index = faiss.IndexHNSWFlat(cfg.embedding_dim, 32) |
|
|
else: |
|
|
raise ValueError(f"Unsupported index_type {cfg.index_type}") |
|
|
self.index.add(embeddings) |
|
|
self.config = cfg |
|
|
|
|
|
def search(self, query_vector: np.ndarray, k: int = 10) -> Tuple[np.ndarray, np.ndarray]: |
|
|
scores, idx = self.index.search(query_vector.astype(np.float32)[None, :], k) |
|
|
return scores[0], idx[0] |
|
|
|
|
|
def save(self, path: str, metadata: dict) -> None: |
|
|
Path(path).parent.mkdir(parents=True, exist_ok=True) |
|
|
faiss.write_index(self.index, path) |
|
|
meta = { |
|
|
**metadata, |
|
|
"index_type": self.config.index_type, |
|
|
"embedding_dim": self.config.embedding_dim, |
|
|
"saved_at": datetime.utcnow().isoformat(), |
|
|
} |
|
|
with open(Path(path).with_suffix(".json"), "w") as f: |
|
|
json.dump(meta, f, indent=2) |
|
|
|
|
|
@classmethod |
|
|
def load(cls, path: str): |
|
|
index = faiss.read_index(path) |
|
|
obj = cls.__new__(cls) |
|
|
obj.index = index |
|
|
obj.config = VectorIndexConfig(index.d) |
|
|
return obj |
|
|
|
