Spaces:
Sleeping
Sleeping
File size: 4,692 Bytes
ca16a03 | 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 | import os
import json
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
class MindGuardVectorDB:
"""
This class handles the ingestion of clinical guidelines (text)
and converts them into mathematical embeddings stored in ChromaDB.
The ingestion pipeline (documents -> embeddings -> vector store) is now
orchestrated entirely through LangChain: each coping strategy becomes a
LangChain `Document`, `HuggingFaceEmbeddings` handles the text-to-vector
conversion, and `Chroma` (from `langchain_chroma`) is the persistent
vector store — a LangChain-native wrapper around the same ChromaDB
collection used before.
"""
def __init__(self):
print("🗄️ Initializing MindGuard Vector Database Builder (LangChain)...")
# --- STRICT ARCHITECTURE PATHING ---
self.script_dir = os.path.dirname(os.path.abspath(__file__))
self.project_root = os.path.abspath(os.path.join(self.script_dir, "../../"))
# Paths aligned perfectly with the folder directory
self.knowledge_base_path = os.path.join(self.project_root, "data", "knowledge_base", "coping_strategies.json")
self.chroma_db_dir = os.path.join(self.project_root, "artifacts", "chroma_db")
# Ensure the Chroma DB output folder exists
os.makedirs(self.chroma_db_dir, exist_ok=True)
# --- LANGCHAIN EMBEDDING ENGINE ---
# Same upgraded model as before (BAAI/bge-base-en-v1.5), now wrapped by
# LangChain's HuggingFaceEmbeddings so the exact same embedding object
# can be reused across the ingestion and retrieval layers.
self.embedding_fn = HuggingFaceEmbeddings(
model_name="BAAI/bge-base-en-v1.5",
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": True},
)
# --- LANGCHAIN VECTOR STORE ---
# `Chroma` from `langchain_chroma` is LangChain's native wrapper around
# a persistent ChromaDB client. It creates/loads the 'clinical_guidelines'
# collection for us.
self.vector_store = Chroma(
collection_name="clinical_guidelines",
embedding_function=self.embedding_fn,
persist_directory=self.chroma_db_dir,
)
print(f"✅ Connected to ChromaDB (via LangChain) at: {self.chroma_db_dir}")
def build_database(self):
"""Reads the JSON file and embeds it into the database as LangChain Documents."""
print(f"📖 Reading clinical data from: {self.knowledge_base_path}...")
# 1. Read the JSON file
with open(self.knowledge_base_path, 'r', encoding='utf-8') as file:
cbt_data = json.load(file)
# 2. Wrap every strategy in a LangChain `Document`
# `page_content` is the text that actually gets embedded and later
# handed to the LLM; `metadata` travels alongside it for filtering.
documents = []
ids = []
for strategy in cbt_data:
documents.append(
Document(
page_content=strategy["content"],
metadata={
"emotion": strategy["primary_emotion"],
"risk_level": strategy["target_risk_level"],
"category": strategy["category"],
"strategy": strategy["strategy_name"],
# Chroma metadata values must be str/int/float/bool, so
# the tag list becomes a single comma-separated string.
"tags": ", ".join(strategy["tags"]),
},
)
)
# Reuse the same generalized ID from the JSON instead of a random one
ids.append(strategy["id"])
print("⚙️ Embedding text into mathematical vectors via LangChain... (This may take a moment to download the model on the first run)")
# 3. Inject into the Vector Database through LangChain
# `add_documents` upserts under the hood, so re-running this script is
# always safe and never creates duplicate entries.
self.vector_store.add_documents(documents=documents, ids=ids)
print(f"✅ Successfully embedded {len(documents)} clinical coping strategies into ChromaDB!")
print("The RAG Knowledge Base is now primed and ready for the Retriever.")
# --- EXECUTION BLOCK ---
if __name__ == "__main__":
db_builder = MindGuardVectorDB()
db_builder.build_database()
|