Update retriever.py

retriever.py

@@ -1,9 +1,89 @@
-from langchain_chroma import Chroma
+#!/usr/bin/env python3
+"""
+Unified retriever with clean reference formatting.
+"""
+
+import os
+from pathlib import Path
 from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
 
-VECTOR_DB_DIR = "vector_db"
 
 def get_retriever():
-    embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en")
-    vectordb = Chroma(persist_directory=VECTOR_DB_DIR, embedding_function=embeddings)
-    return vectordb.as_retriever(search_kwargs={"k": 3})
+    """Get retriever with consistent formatting"""
+
+    # Locate vector_db
+    vector_db_path = Path("vector_db")
+    if not vector_db_path.exists():
+        vector_db_path = Path("../vector_db")
+
+    if not vector_db_path.exists():
+        raise FileNotFoundError("Vector database not found")
+
+    print(f"Loading from: {vector_db_path.absolute()}")
+
+    # Init embedding model
+    embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en")
+
+    # Load the correct collection
+    vectordb = Chroma(
+        persist_directory=str(vector_db_path.absolute()),
+        embedding_function=embedding_model,
+        collection_name="legal_documents"
+    )
+
+    count = vectordb._collection.count()
+    print(f"Found {count} documents in collection 'legal_documents'")
+
+    if count == 0:
+        raise ValueError("Collection 'legal_documents' is empty")
+
+    return vectordb.as_retriever(search_kwargs={"k": 5})
+
+
+def format_reference(doc) -> str:
+    """
+    Format metadata into: [doc_type] Section — Source
+    """
+    doc_type = doc.metadata.get("document_type", "general")
+    section = doc.metadata.get("section", "Unknown Section")
+    source = doc.metadata.get("source", "Unknown Source")
+    return f"[{doc_type}] {section} — {source}"
+
+
+def test_retriever():
+    """Test the retriever with sample queries"""
+    try:
+        retriever = get_retriever()
+
+        test_queries = [
+            "constitutional rights",
+            "FCCPA tribunal",
+            "labour law employment",
+            "data protection"
+        ]
+
+        for query in test_queries:
+            print(f"\nTesting: '{query}'")
+            results = retriever.invoke(query)
+            print(f"  Found {len(results)} results")
+
+            for doc in results:
+                ref = format_reference(doc)
+                preview = doc.page_content[:150].replace("\n", " ")
+                print(f"  {ref}")
+                print(f"    Preview: {preview}...")
+
+        return True
+
+    except Exception as e:
+        print(f"✗ Error: {e}")
+        return False
+
+
+if __name__ == "__main__":
+    success = test_retriever()
+    if success:
+        print("\n✓ Retriever is working correctly with formatted references!")
+    else:
+        print("\n✗ Retriever test failed")