Spaces:
Sleeping
Sleeping
File size: 4,870 Bytes
7a11b03 | 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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """Embedding and ChromaDB vector-store functions."""
from functools import lru_cache
import hashlib
import os
os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")
import chromadb
from sentence_transformers import SentenceTransformer
from src.config import CHROMA_DIR, COLLECTION_NAME, EMBEDDING_MODEL_NAME, TOP_K_RESULTS
@lru_cache(maxsize=1)
def load_embedding_model() -> SentenceTransformer:
"""Load the embedding model once so repeated questions are faster."""
return SentenceTransformer(EMBEDDING_MODEL_NAME)
@lru_cache(maxsize=1)
def get_chroma_client():
"""Create one persistent ChromaDB client for the running server."""
CHROMA_DIR.mkdir(parents=True, exist_ok=True)
return chromadb.PersistentClient(path=str(CHROMA_DIR))
def get_chroma_collection():
"""Create or open the ChromaDB collection used by this app."""
client = get_chroma_client()
return client.get_or_create_collection(name=COLLECTION_NAME)
def clear_vector_database() -> None:
"""Clear old notes by resetting the ChromaDB collection."""
client = get_chroma_client()
try:
client.delete_collection(name=COLLECTION_NAME)
except Exception:
pass
client.get_or_create_collection(name=COLLECTION_NAME)
def get_vector_count() -> int:
"""Return how many chunks are stored in ChromaDB."""
collection = get_chroma_collection()
return collection.count()
def store_chunks_in_vector_database(chunks: list[dict]) -> int:
"""Convert chunks into embeddings and save them in ChromaDB."""
if not chunks:
raise ValueError("There are no chunks to store in the vector database.")
model = load_embedding_model()
collection = get_chroma_collection()
texts = [chunk["text"] for chunk in chunks]
embeddings = model.encode(texts, show_progress_bar=False).tolist()
ids = [
hashlib.sha1(
f"{chunk['source_file']}|{chunk['page_number']}|{chunk['chunk_id']}|{chunk['text']}".encode("utf-8")
).hexdigest()
for chunk in chunks
]
metadata = [
{
"chunk_id": chunk["chunk_id"],
"page_number": chunk["page_number"],
"source_file": chunk["source_file"],
"extraction_method": chunk.get("extraction_method", "embedded_text"),
"ocr_mode": chunk.get("ocr_mode", "not_used"),
"ocr_engine": chunk.get("ocr_engine", "not_used"),
}
for chunk in chunks
]
collection.add(ids=ids, documents=texts, embeddings=embeddings, metadatas=metadata)
return len(chunks)
def search_relevant_chunks(question: str, top_k: int = TOP_K_RESULTS) -> list[dict]:
"""Find chunks that are most similar to the user's question."""
if not question.strip():
raise ValueError("Please enter a question first.")
model = load_embedding_model()
collection = get_chroma_collection()
if collection.count() == 0:
raise ValueError("Please upload and process notes before asking a question.")
question_embedding = model.encode([question], show_progress_bar=False).tolist()[0]
results = collection.query(
query_embeddings=[question_embedding],
n_results=top_k,
include=["documents", "metadatas", "distances"],
)
relevant_chunks = []
documents = results.get("documents", [[]])[0]
metadatas = results.get("metadatas", [[]])[0]
distances = results.get("distances", [[]])[0]
for document, metadata, distance in zip(documents, metadatas, distances):
relevant_chunks.append(
{
"text": document,
"page_number": metadata["page_number"],
"source_file": metadata["source_file"],
"chunk_id": metadata["chunk_id"],
"extraction_method": metadata.get("extraction_method", "embedded_text"),
"ocr_mode": metadata.get("ocr_mode", "not_used"),
"ocr_engine": metadata.get("ocr_engine", "not_used"),
"distance": distance,
}
)
return relevant_chunks
def get_all_stored_chunks(max_characters: int = 12000) -> str:
"""Read stored chunks and combine them for summary-style tasks."""
collection = get_chroma_collection()
if collection.count() == 0:
raise ValueError("Please upload and process notes first.")
results = collection.get(include=["documents", "metadatas"])
combined_parts = []
current_length = 0
for document, metadata in zip(results["documents"], results["metadatas"]):
source_label = f"[{metadata['source_file']} - Page {metadata['page_number']}]"
part = f"{source_label}\n{document}"
if current_length + len(part) > max_characters:
break
combined_parts.append(part)
current_length += len(part)
return "\n\n".join(combined_parts)
|