Updated

smebuilder_vector.py

@@ -1,24 +1,57 @@
 import os
+import pandas as pd
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Chroma
+from langchain.schema import Document
 
-# Make sure HF_HOME is set for cache
-os.environ["HF_HOME"] = os.getenv("HF_HOME", "/app/huggingface_cache")
-os.environ["TRANSFORMERS_CACHE"] = os.environ["HF_HOME"]
+# ----------------- CONFIG -----------------
+DATASET_PATH = "sme_builder_dataset.csv"
+DB_LOCATION = os.getenv("CHROMA_DB_DIR", "./Dev_Assist_SME_Builder_DB")
+COLLECTION_NAME = "landing_page_generation_examples"
+EMBEDDING_MODEL = os.getenv("HF_EMBEDDING_MODEL", "intfloat/e5-large-v2")
+HF_CACHE_DIR = os.getenv("HF_CACHE_DIR", "/app/huggingface_cache")
 
-# Initialize embeddings (without cache_dir argument)
+os.makedirs(HF_CACHE_DIR, exist_ok=True)
+os.makedirs(DB_LOCATION, exist_ok=True)
+
+# ----------------- LOAD DATASET -----------------
+if not os.path.exists(DATASET_PATH):
+    raise FileNotFoundError(f"Dataset file not found: {DATASET_PATH}")
+
+df = pd.read_csv(DATASET_PATH)
+
+# ----------------- EMBEDDINGS -----------------
 embeddings = HuggingFaceEmbeddings(
-    model_name="sentence-transformers/all-MiniLM-L6-v2",
-    model_kwargs={"device": "cpu"}  # or "cuda" if GPU available
+    model_name=EMBEDDING_MODEL,
+    cache_dir=HF_CACHE_DIR
 )
 
-# Chroma vector store (writable directory)
-chroma_db_path = os.getenv("CHROMA_DB_DIR", "/app/Dev_Assist_SME_Builder_DB")
+# ----------------- VECTOR STORE -----------------
+# Only add documents if DB is empty
+add_documents = not os.listdir(DB_LOCATION)
+
 vector_store = Chroma(
-    persist_directory=chroma_db_path,
+    collection_name=COLLECTION_NAME,
+    persist_directory=DB_LOCATION,
     embedding_function=embeddings,
-    collection_name="sme_collection"
 )
 
-# retriever object for querying
-retriever = vector_store.as_retriever()
+if add_documents:
+    documents = []
+    for i, row in df.iterrows():
+        content = " ".join([
+            str(row.get("prompt", "")),
+            str(row.get("html_code", "")),
+            str(row.get("css_code", "")),
+            str(row.get("js_code", "")),
+            str(row.get("sector", ""))
+        ]).strip()
+        documents.append(Document(page_content=content, id=str(i)))
+    
+    if documents:
+        vector_store.add_documents(documents=documents, ids=[doc.id for doc in documents])
+
+# ----------------- RETRIEVER -----------------
+retriever = vector_store.as_retriever(search_kwargs={"k": 20})
+
+print(f"Vector store ready with {len(df)} documents.")