Spaces:
Running
Running
File size: 1,907 Bytes
e43c82d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """Vector Store implementation for indexing and semantic similarity search.
A vector database stores chunks alongside their embedding vectors and finds
the nearest neighbors (most relevant chunks) using cosine similarity.
"""
from rag_engine.embeddings import cosine_similarity, get_embedding_model
class VectorStore:
"""Simple in-memory vector store with cosine similarity ranking."""
def __init__(self, embedding_model=None):
self.embedding_model = embedding_model or get_embedding_model()
self.chunks = []
self.vectors = []
def add_chunks(self, chunks):
"""Add TextChunk objects and calculate their embeddings."""
if not chunks:
return
new_texts = [c.text for c in chunks]
new_vectors = self.embedding_model.embed_documents(new_texts)
self.chunks.extend(chunks)
self.vectors.extend(new_vectors)
def search(self, query, top_k=5):
"""Find the top_k most semantically similar chunks for a query.
Returns:
List of dicts: [{"chunk": TextChunk, "score": float, "source": "dense"}]
"""
if not self.chunks:
return []
query_vector = self.embedding_model.embed_query(query)
scored_results = []
for idx, (chunk, doc_vector) in enumerate(zip(self.chunks, self.vectors)):
similarity = cosine_similarity(query_vector, doc_vector)
scored_results.append({
"chunk": chunk,
"score": round(similarity, 4),
"source": "dense",
})
# Sort descending by similarity score
scored_results.sort(key=lambda x: x["score"], reverse=True)
return scored_results[:top_k]
def clear(self):
"""Empty the store."""
self.chunks = []
self.vectors = []
def __len__(self):
return len(self.chunks)
|