Spaces:
Running
Running
File size: 3,503 Bytes
74e8a7b | 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 | """
ChromaDB collection access and similarity/metadata search.
"""
import chromadb
from src.utils.config import DB_PATH, COLLECTION_NAME, SIMILARITY_THRESHOLD
from src.knowledge_base.embeddings import embed_texts_with_retry
_client = chromadb.PersistentClient(path=DB_PATH)
collection = _client.get_collection(COLLECTION_NAME)
def search(
query: str,
top_k: int = 3,
metadata_filter: dict | None = None,
where: dict | None = None,
is_exact_fetch: bool = False,
) -> dict:
"""
Embed query and search ChromaDB.
metadata_filter: optional 'where' clause for similarity search.
where: 'where' clause for a direct metadata lookup (used when is_exact_fetch=True).
is_exact_fetch: skip embedding similarity entirely and fetch ALL chunks matching `where`
via collection.get(). Used for structured category lookups
(study_plan / program_info / scholarship / all_programs comparison)
where the whole matching set is wanted, not a top-k similarity guess.
NOTE: matches current notebook behavior — has_answer is returned True even when
documents is empty on the exact-fetch path (only a warning is printed). This means
the fallback-to-similarity-search in multi_query_search never actually triggers on
a truly empty exact-fetch result. Recommend restoring an `if documents:` guard here
if you want the fallback to work as originally intended — flagging for your decision.
"""
# ── Study plan: direct metadata lookup, no ranking ─────────────────
if is_exact_fetch and where:
results = collection.get(where=where)
documents = results["documents"]
metadatas = results["metadatas"]
if not documents:
print(" [search] No exact-category chunks matched — falling back to similarity search")
scores = [1.0] * len(documents)
return {"has_answer": True, "documents": documents, "metadatas": metadatas,
"scores": scores, "best_score": 1.0}
# ── Embedding-based similarity search ──────────────────────────────
query_embedding = embed_texts_with_retry([query])[0]
query_params = dict(
query_embeddings=[query_embedding],
n_results=top_k,
include=["documents", "metadatas", "distances"],
)
if metadata_filter:
query_params["where"] = metadata_filter
results = collection.query(**query_params)
documents = results["documents"][0]
metadatas = results["metadatas"][0]
distances = results["distances"][0]
# Fallback: if filter matched nothing, retry without the filter
if not documents and metadata_filter:
query_params.pop("where", None)
results = collection.query(**query_params)
documents = results["documents"][0]
metadatas = results["metadatas"][0]
distances = results["distances"][0]
if not documents:
return {"has_answer": False, "documents": [], "metadatas": [], "scores": [], "best_score": 0.0}
scores = [1 - d for d in distances]
best_score = max(scores)
if best_score < SIMILARITY_THRESHOLD:
return {"has_answer": False, "documents": [], "metadatas": [], "scores": [], "best_score": best_score}
return {"has_answer": True, "documents": documents, "metadatas": metadatas, "scores": scores, "best_score": best_score}
|