Update app.py

app.py

@@ -2,17 +2,29 @@ import os
 from pathlib import Path
 import gradio as gr
 
-from retriever import get_retriever
+
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_community.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
+from langchain_community.vectorstores import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
 
-# Ensure vector DB exists (from complete_ingestion.py output)
+# Load the vector DB created by complete_ingestion.py
 PERSIST_DIR = Path("data/processed/vector_db")
+
 if not PERSIST_DIR.exists() or not any(PERSIST_DIR.iterdir()):
-    raise RuntimeError("⚠️ Vector DB not found. Please run complete_ingestion.py first.")
+    print("⚠️ Vector DB not found. Run complete_ingestion.py first.")
+    raise SystemExit(1)
+
+embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en")
+vectordb = Chroma(
+    persist_directory=str(PERSIST_DIR),
+    embedding_function=embedding_model,
+    collection_name="legal_documents"
+)
+
+retriever = vectordb.as_retriever(search_kwargs={"k": 3})
 
-retriever = get_retriever()
 
 # Load lightweight conversational model
 MODEL_ID = os.getenv("LLM_ID", "TinyLlama/TinyLlama-1.1B-Chat-v1.0")