feat: rebase and implement file API methods

Signed-off-by: Varsha Prasad Narsing <varshaprasad96@gmail.com>
This commit is contained in:
Varsha Prasad Narsing 2025-06-25 16:59:29 -07:00
parent 61bddfe70e
commit 3d27b7054c
13 changed files with 195 additions and 141 deletions

View file

@ -214,7 +214,7 @@ def _validate_embedding(embedding: NDArray, index: int, expected_dimension: int)
class EmbeddingIndex(ABC):
@abstractmethod
async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray):
async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray, metadata: dict[str, Any] | None = None):
raise NotImplementedError()
@abstractmethod
@ -251,6 +251,7 @@ class VectorDBWithIndex:
async def insert_chunks(
self,
chunks: list[Chunk],
distance_metric: str | None = None,
) -> None:
chunks_to_embed = []
for i, c in enumerate(chunks):
@ -271,7 +272,13 @@ class VectorDBWithIndex:
c.embedding = embedding
embeddings = np.array([c.embedding for c in chunks], dtype=np.float32)
await self.index.add_chunks(chunks, embeddings)
# Create metadata dict with distance_metric if provided
metadata = None
if distance_metric is not None:
metadata = {"distance_metric": distance_metric}
await self.index.add_chunks(chunks, embeddings, metadata=metadata)
async def query_chunks(
self,