Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from pathlib import Path | |
| from typing import Any | |
| import chromadb | |
| from src.schemas.domain import ChunkRecord | |
| class PaperVectorStore: | |
| def __init__(self, directory: Path, collection_name: str = "papermind") -> None: | |
| directory.mkdir(parents=True, exist_ok=True) | |
| self.client = chromadb.PersistentClient(path=str(directory)) | |
| self.collection_name = collection_name | |
| self.collection = self.client.get_or_create_collection( | |
| name=collection_name, | |
| configuration={ | |
| "hnsw": { | |
| "space": "cosine", | |
| "ef_construction": 200, | |
| "ef_search": 160, | |
| } | |
| }, | |
| metadata={"description": "PaperMind session evidence index"}, | |
| ) | |
| def reset(self) -> None: | |
| try: | |
| self.client.delete_collection(self.collection_name) | |
| except Exception: | |
| pass | |
| self.collection = self.client.get_or_create_collection( | |
| name=self.collection_name, | |
| configuration={"hnsw": {"space": "cosine", "ef_search": 160}}, | |
| metadata={"description": "PaperMind session evidence index"}, | |
| ) | |
| def add(self, chunks: list[ChunkRecord], embeddings: list[list[float]]) -> None: | |
| if len(chunks) != len(embeddings): | |
| raise ValueError("Chunk and embedding counts do not match.") | |
| if not chunks: | |
| return | |
| self.collection.add( | |
| ids=[chunk.chunk_id for chunk in chunks], | |
| documents=[chunk.text for chunk in chunks], | |
| metadatas=[chunk.metadata() for chunk in chunks], | |
| embeddings=embeddings, | |
| ) | |
| def dense_query(self, embedding: list[float], limit: int) -> list[dict[str, Any]]: | |
| result = self.collection.query( | |
| query_embeddings=[embedding], | |
| n_results=min(limit, self.count()), | |
| include=["documents", "metadatas", "distances"], | |
| ) | |
| if not result.get("ids") or not result["ids"][0]: | |
| return [] | |
| rows: list[dict[str, Any]] = [] | |
| for chunk_id, document, metadata, distance in zip( | |
| result["ids"][0], | |
| result["documents"][0], | |
| result["metadatas"][0], | |
| result["distances"][0], | |
| strict=True, | |
| ): | |
| rows.append( | |
| { | |
| "chunk_id": chunk_id, | |
| "document": document, | |
| "metadata": metadata, | |
| "distance": float(distance), | |
| "score": max(0.0, 1.0 - float(distance)), | |
| } | |
| ) | |
| return rows | |
| def all_chunks(self, include_embeddings: bool = False) -> list[dict[str, Any]]: | |
| include = ["documents", "metadatas"] | |
| if include_embeddings: | |
| include.append("embeddings") | |
| result = self.collection.get(include=include) | |
| rows: list[dict[str, Any]] = [] | |
| embeddings = result.get("embeddings") | |
| for index, chunk_id in enumerate(result.get("ids", [])): | |
| row = { | |
| "chunk_id": chunk_id, | |
| "document": result["documents"][index], | |
| "metadata": result["metadatas"][index], | |
| } | |
| if embeddings is not None: | |
| row["embedding"] = embeddings[index] | |
| rows.append(row) | |
| return rows | |
| def fetch_embeddings(self, chunk_ids: list[str]) -> dict[str, list[float]]: | |
| if not chunk_ids: | |
| return {} | |
| result = self.collection.get(ids=chunk_ids, include=["embeddings"]) | |
| embeddings = result.get("embeddings") | |
| if embeddings is None: | |
| return {} | |
| return { | |
| chunk_id: list(vector) | |
| for chunk_id, vector in zip(result["ids"], embeddings, strict=True) | |
| } | |
| def count(self) -> int: | |
| return self.collection.count() | |