Project: a help-centre assistant with citations and permissions

import os

from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

from extractive_llm import ExtractiveLLM

Settings.embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
Settings.llm = ExtractiveLLM()

def with_audience(path):
    audience = "staff" if "/staff/" in path else "customer"
    return {"file_name": os.path.basename(path), "audience": audience}


documents = SimpleDirectoryReader(".", recursive=True, required_exts=[".md"], file_metadata=with_audience).load_data()
for document in documents:
    document.excluded_embed_metadata_keys = ["audience"]
    document.excluded_llm_metadata_keys = ["audience"]
index = VectorStoreIndex.from_documents(documents)

ALLOWED = {"customer": ["customer"], "staff": ["customer", "staff"]}


def ask(question, role):
    filters = MetadataFilters(
        filters=[ExactMatchFilter(key="audience", value=a) for a in ALLOWED[role]], condition="or"
    )
    engine = index.as_query_engine(
        similarity_top_k=2,
        filters=filters,
        node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.3)],
    )
    response = engine.query(question)
    if not response.source_nodes:
        return "I can't find that in our help centre. A person from the team will reply."
    sources = sorted({n.metadata["file_name"] for n in response.source_nodes})
    return f"{response} (sources: {', '.join(sources)})"

for question, role in [
    ("How long until my refund money reaches my card?", "customer"),
    ("Who approves a refund over 200?", "customer"),
    ("Who approves a refund over 200?", "staff"),
    ("What are your opening hours?", "customer"),
]:
    print(f"[{role}] {question}")
    print("   ", ask(question, role))

python assistant.py