codescraftman / ethanicbot / langchain / retriever.py
retriever.py
Raw
import os
from django.apps import apps
from langchain.schema import Document
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

# Use a lightweight and effective embedding model
EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
VECTORSTORE_PATH = "ethanicbot/langchain/faiss_index"

def load_all_model_data():
    """
    Load data from all models in the Django project and create a list of Document objects.
    Each document contains the model's field data and associated metadata.
    """
    docs = []

    for model in apps.get_models():
        for instance in model.objects.all():
            model_name = model.__name__
            app_label = model._meta.app_label
            metadata = {
                "model": model_name,
                "app": app_label,
                "id": instance.pk,
            }

            fields = []
            for field in model._meta.fields:
                try:
                    value = getattr(instance, field.name)
                    fields.append(f"{field.name}: {value}")
                except Exception as e:
                    print(f"⚠️ Error getting field {field.name} for instance {instance.pk}: {e}")
                    continue

            content = f"Model: {model_name} | App: {app_label}\n" + "\n".join(fields)
            docs.append(Document(page_content=content, metadata=metadata))

    return docs

def create_vectorstore():
    """
    Create a FAISS vector store from all model data, and save it to the specified path.
    """
    try:
        docs = load_all_model_data()

        if not docs:
            raise ValueError("No documents to index. Ensure models have data.")

        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
        vectorstore = FAISS.from_documents(docs, embeddings)
        vectorstore.save_local(VECTORSTORE_PATH)
        print("✅ Vectorstore created and saved successfully.")
    except Exception as e:
        print(f"❌ Error creating vectorstore: {e}")

def load_vectorstore():
    """
    Load the FAISS vector store from the saved location.
    """
    try:
        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
        vectorstore = FAISS.load_local(VECTORSTORE_PATH, embeddings, allow_dangerous_deserialization=True)
        print("✅ Vectorstore loaded successfully.")
        return vectorstore
    except Exception as e:
        print(f"❌ Error loading vectorstore: {e}")
        return None