| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| class VectorStore: |
| """An in-memory semantic index over text chunks, each tagged with source metadata for citations.""" |
|
|
| def __init__(self): |
| self.chunks: list[str] = [] |
| self.metadatas: list[dict] = [] |
| self.embeddings: np.ndarray | None = None |
|
|
| def add(self, chunks: list[str], metadatas: list[dict], embeddings: np.ndarray) -> None: |
| self.chunks.extend(chunks) |
| self.metadatas.extend(metadatas) |
| if self.embeddings is None: |
| self.embeddings = embeddings |
| else: |
| self.embeddings = np.vstack([self.embeddings, embeddings]) |
|
|
| def search(self, query_embedding: np.ndarray, top_k: int = 5) -> list[dict]: |
| """Return the top_k most similar chunks, each with its source metadata and similarity score.""" |
| scores = np.dot(self.embeddings, query_embedding) |
| top_indices = np.argsort(scores)[-top_k:][::-1] |
| return [ |
| {"text": self.chunks[i], "score": float(scores[i]), **self.metadatas[i]} |
| for i in top_indices |
| ] |
|
|
| def save(self, dir_path: str) -> None: |
| """Persist the index to disk so it doesn't need to be rebuilt on every run.""" |
| dest = Path(dir_path) |
| dest.mkdir(parents=True, exist_ok=True) |
| np.save(dest / "embeddings.npy", self.embeddings) |
| records = [{"text": chunk, **metadata} for chunk, metadata in zip(self.chunks, self.metadatas)] |
| (dest / "records.json").write_text(json.dumps(records, indent=2)) |
|
|
| @classmethod |
| def load(cls, dir_path: str) -> "VectorStore": |
| dest = Path(dir_path) |
| records = json.loads((dest / "records.json").read_text()) |
| store = cls() |
| store.chunks = [record["text"] for record in records] |
| store.metadatas = [{k: v for k, v in record.items() if k != "text"} for record in records] |
| store.embeddings = np.load(dest / "embeddings.npy") |
| return store |
|
|
