Spaces:
Sleeping
Sleeping
File size: 5,888 Bytes
fe2ee65 | 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 | import json
import os
from openai import OpenAI
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
class NvidiaCompatibleEmbeddings(Embeddings):
"""
Custom embedding processor that bypasses LangChain's internal tokenization middleware
to pass raw text strings and asymmetric type configurations directly to NVIDIA NIM.
"""
def __init__(self, model: str, api_key: str, base_url: str):
self.client = OpenAI(api_key=api_key, base_url=base_url)
self.model = model
def embed_documents(self, texts: list[str]) -> list[list[float]]:
"""Embeds a list of documentation chunks using the 'passage' type."""
response = self.client.embeddings.create(
input=texts,
model=self.model,
extra_body={"input_type": "passage"}
)
return [item.embedding for item in response.data]
def embed_query(self, text: str) -> list[float]:
"""Embeds a live user search query using the 'query' type."""
response = self.client.embeddings.create(
input=[text],
model=self.model,
extra_body={"input_type": "query"}
)
return response.data[0].embedding
class VectorDBManager:
def __init__(self, persist_directory: str = "./chroma_db"):
self.embeddings = NvidiaCompatibleEmbeddings(
model=os.getenv("EMBEDDING_MODEL", "nvidia/llama-nemotron-embed-1b-v2"),
api_key=os.getenv("NVIDIA_API_KEY"),
base_url=os.getenv("NVIDIA_BASE_URL")
)
self.persist_directory = persist_directory
self.vector_store = None
def _ensure_vector_store(self):
"""🛡️ Internal safeguard to ensure the Chroma instance is actively loaded in memory."""
if self.vector_store is None:
self.vector_store = Chroma(
persist_directory=self.persist_directory,
embedding_function=self.embeddings
)
def initialize_db(self, faq_filepath: str):
"""
Loads incoming JSON FAQs, formats them into searchable LangChain Documents,
and saves them locally via ChromaDB stamped with a 'global' scope.
"""
if not os.path.exists(faq_filepath):
raise FileNotFoundError(f"Could not find FAQ data resource file at: {faq_filepath}")
with open(faq_filepath, 'r') as f:
faq_data = json.load(f)
documents = []
for item in faq_data:
page_content = f"Question: {item['question']}\nAnswer: {item['answer']}"
# 🔑 Added session_id: "global" so these base FAQs are accessible to all users
metadata = {
"category": item["category"],
"faq_id": item["id"],
"session_id": "global"
}
documents.append(Document(page_content=page_content, metadata=metadata))
self.vector_store = Chroma.from_documents(
documents=documents,
embedding=self.embeddings,
persist_directory=self.persist_directory
)
print(f"🚀 Vector DB successfully initialized with {len(documents)} FAQs via NVIDIA Embeddings.")
def get_retriever(self, session_id: str = "default_session"):
"""
Loads the localized database and converts it into a queryable retriever layer
strictly filtered by the active session identity.
"""
self._ensure_vector_store()
# 🔑 MULTI-TENANCY FILTER: Look for global baseline data OR this specific user's uploads
meta_filter = {
"$or": [
{"session_id": "global"},
{"session_id": session_id}
]
}
return self.vector_store.as_retriever(
search_kwargs={
"k": 2,
"filter": meta_filter
}
)
def add_text_to_db(self, text: str, filename: str, session_id: str = "default_session"):
"""
📥 Chunks raw text from dynamic frontend uploads, computes NVIDIA embeddings,
and appends them directly into the live database marked with the owner's session ID.
"""
self._ensure_vector_store()
# 1. Break the document down into semantic pieces
text_splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=60)
chunks = text_splitter.split_text(text)
# 2. Package raw strings into formal LangChain Document structures
documents = []
for idx, chunk in enumerate(chunks):
# 🔑 Stamping chunk dictionary metadata with the active session tracking ID
metadata = {
"source": filename,
"category": "Dynamic Upload",
"chunk_index": idx,
"session_id": session_id
}
documents.append(Document(page_content=chunk, metadata=metadata))
# 3. Stream the newly generated embeddings straight into the persistent store
self.vector_store.add_documents(documents)
print(f"⚡ Successfully indexed {len(documents)} dynamic chunks from raw file: '{filename}' for session '{session_id}'")
def clear_db(self):
"""🗑️ Completely wipes out the existing vector store collection securely on disk and in RAM."""
try:
self._ensure_vector_store()
self.vector_store.delete_collection()
self.vector_store = None
print("🗑️ Vector database collection successfully cleared!")
return True
except Exception as e:
print(f"❌ Failed to clear vector database: {e}")
return False |