How to Build an In-Memory Vector Store in Python

import math
from typing import Dict, List, Optional


class InMemoryVectorStore:
    def __init__(self) -> None:
        self.vectors: Dict[str, List[float]] = {}
        self.index: Dict[str, List[str]] = {}  # query -> list of ids sorted by similarity

    def add(self, vector_id: str, vector: List[float]) -> None:
        self.vectors[vector_id] = vector
        self.index.clear()  # invalidate cache on any insertion

    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))
        if norm_a == 0.0 or norm_b == 0.0:
            return 0.0
        return dot / (norm_a * norm_b)

    def search(self, query_vector: List[float], top_k: int = 3) -> List[tuple]:
        if not self.vectors:
            return []
        scored = [(vector_id, self._cosine_similarity(query_vector, vec))
                  for vector_id, vec in self.vectors.items()]
        scored.sort(key=lambda item: item[1], reverse=True)
        return scored[:top_k]


if __name__ == "__main__":
    store = InMemoryVectorStore()
    store.add("doc1", [1.0, 0.0, 0.0])
    store.add("doc2", [0.8, 0.6, 0.0])
    store.add("doc3", [0.0, 1.0, 0.0])

    results = store.search([1.0, 0.0, 0.0], top_k=2)
    print(results)