Spaces:
Sleeping
Sleeping
File size: 3,432 Bytes
ca20ec1 | 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 | import json
import os
import chromadb
class VectorDBManager:
def __init__(self, collection_name: str = "shl_catalog"):
# Persistent storage locally, fits stateless server startups perfectly
self.client = chromadb.Client()
self.collection = self.client.get_or_create_collection(name=collection_name)
def initialize_catalog(self, catalog_path: str):
"""Loads and indexes the catalog JSON file if the collection is empty."""
if self.collection.count() > 0:
return # Already populated
if not os.path.exists(catalog_path):
raise FileNotFoundError(f"Catalog file not found at: {catalog_path}")
# Added utf-8 encoding and strict=False to bypass dirty JSON characters
with open(catalog_path, "r", encoding="utf-8") as f:
catalog = json.load(f, strict=False)
docs = []
metadatas = []
ids = []
for item in catalog:
# Combine content for meaningful semantic retrieval
text_chunk = (
f"Name: {item['name']}. "
f"Description: {item['description']}. "
f"Job Levels: {', '.join(item.get('job_levels', []))}"
)
docs.append(text_chunk)
# Metadata aligns precisely with the output Recommendation model
metadatas.append({
"name": item["name"],
"url": item.get("link", ""),
"test_type": item.get("keys", ["General"])[0] if item.get("keys") else "General"
})
ids.append(str(item["entity_id"]))
self.collection.add(documents=docs, metadatas=metadatas, ids=ids)
def query_catalog(self, query_text: str, n_results: int = 8) -> str:
"""Queries the collection and returns a structured string for prompt context.
Kept for backward compatibility — prefer query_catalog_structured for new code,
since the structured form is what enables post-generation validation."""
results = self.collection.query(query_texts=[query_text], n_results=n_results)
context_str = ""
if results['documents'] and results['documents'][0]:
for i, doc in enumerate(results['documents'][0]):
meta = results['metadatas'][0][i]
context_str += f"- {doc} (URL: {meta['url']}, Type: {meta['test_type']})\n"
return context_str
def query_catalog_structured(self, query_text: str, n_results: int = 8) -> list:
"""Same query, returned as a list of dicts (name/url/test_type/description)
instead of a flattened string. This is the ground truth used both to build
the responder's candidate list AND to validate its output afterward —
anything the model returns that isn't in this list gets dropped."""
results = self.collection.query(query_texts=[query_text], n_results=n_results)
candidates = []
if results['documents'] and results['documents'][0]:
for i, doc in enumerate(results['documents'][0]):
meta = results['metadatas'][0][i]
candidates.append({
"name": meta["name"],
"url": meta["url"],
"test_type": meta["test_type"],
"description": doc,
})
return candidates
# Global singleton instance
db_manager = VectorDBManager() |