query_understanding.py

import re
import spacy
from spacy.pipeline import EntityRuler

from transformers import pipeline
from sentence_transformers import SentenceTransformer, util


# -----------------------------
# Load Models
# -----------------------------

print("Loading models...")


# Intent classifier
intent_classifier = pipeline(
    "zero-shot-classification",
    model="facebook/bart-large-mnli"
)


# Query rewriting model
# NOTE: flan-t5 is a seq2seq model, so it needs "text2text-generation",
# not "text-generation" (that pipeline expects a decoder-only/causal LM
# and will error or behave incorrectly with flan-t5).
rewriter = pipeline(
    "text2text-generation",
    model="google/flan-t5-base"
)


# Embedding model
embedding_model = SentenceTransformer(
    "all-MiniLM-L6-v2"
)


# NER model
nlp = spacy.load(
    "en_core_web_sm"
)


# Entity Ruler patterns — registered once, outside the function.
# Adding a pipe named "entity_ruler" on every call to extract_entities()
# would raise "already exists in pipeline" from the second call onward.
ENTITY_PATTERNS = [
    # =======================
    # AI Concepts
    # =======================
    {"label": "AI_CONCEPT", "pattern": "RAG"},
    {"label": "AI_CONCEPT", "pattern": "Retrieval Augmented Generation"},
    {"label": "AI_CONCEPT", "pattern": "LLM"},
    {"label": "AI_CONCEPT", "pattern": "Generative AI"},
    {"label": "AI_CONCEPT", "pattern": "Machine Learning"},
    {"label": "AI_CONCEPT", "pattern": "Deep Learning"},

    # =======================
    # AI Models
    # =======================
    {"label": "AI_MODEL", "pattern": "GPT-4"},
    {"label": "AI_MODEL", "pattern": "GPT-4o"},
    {"label": "AI_MODEL", "pattern": "GPT-5"},
    {"label": "AI_MODEL", "pattern": "Llama 2"},
    {"label": "AI_MODEL", "pattern": "Llama 3"},
    {"label": "AI_MODEL", "pattern": "Claude"},
    {"label": "AI_MODEL", "pattern": "Gemini"},
    {"label": "AI_MODEL", "pattern": "Mistral"},
    {"label": "AI_MODEL", "pattern": "DeepSeek"},

    # =======================
    # Frameworks
    # =======================
    {"label": "FRAMEWORK", "pattern": "LangChain"},
    {"label": "FRAMEWORK", "pattern": "LlamaIndex"},
    {"label": "FRAMEWORK", "pattern": "Haystack"},
    {"label": "FRAMEWORK", "pattern": "CrewAI"},
    {"label": "FRAMEWORK", "pattern": "LangGraph"},

    # =======================
    # Vector Databases
    # =======================
    {"label": "VECTOR_DATABASE", "pattern": "FAISS"},
    {"label": "VECTOR_DATABASE", "pattern": "Chroma"},
    {"label": "VECTOR_DATABASE", "pattern": "Pinecone"},
    {"label": "VECTOR_DATABASE", "pattern": "Weaviate"},
    {"label": "VECTOR_DATABASE", "pattern": "Milvus"},
    {"label": "VECTOR_DATABASE", "pattern": "Qdrant"},

    # =======================
    # Databases
    # =======================
    {"label": "DATABASE", "pattern": "MongoDB"},
    {"label": "DATABASE", "pattern": "MySQL"},
    {"label": "DATABASE", "pattern": "PostgreSQL"},
    {"label": "DATABASE", "pattern": "SQLite"},

    # =======================
    # Programming Languages
    # =======================
    {"label": "LANGUAGE", "pattern": "Python"},
    {"label": "LANGUAGE", "pattern": "Java"},
    {"label": "LANGUAGE", "pattern": "C++"},
    {"label": "LANGUAGE", "pattern": "JavaScript"},

    # =======================
    # Cloud Platforms
    # =======================
    {"label": "CLOUD", "pattern": "AWS"},
    {"label": "CLOUD", "pattern": "Azure"},
    {"label": "CLOUD", "pattern": "Google Cloud"},
]

_ruler = nlp.add_pipe("entity_ruler", before="ner")
_ruler.add_patterns(ENTITY_PATTERNS)


# -----------------------------
# 1. Query Preprocessing
# -----------------------------

def preprocess(query):

    query = query.lower()

    query = re.sub(
        r"[^a-zA-Z0-9 ]",
        "",
        query
    )

    return query.strip()



# -----------------------------
# 2. Intent Classification
# -----------------------------
#
# Only 3 downstream agents exist (Document, SQL, Web Search), so the
# zero-shot labels map directly onto them instead of a larger intent
# taxonomy that then needs re-bucketing at routing time.

INTENT_LABELS = [
    "document lookup",       # -> Document Agent (RAG over stored docs)
    "database query",        # -> SQL Agent (structured/transactional data)
    "web search"              # -> Web Search Agent (fresh/external info)
]

def classify_intent(query):

    result = intent_classifier(
        query,
        INTENT_LABELS
    )

    return {
        "intent": result["labels"][0],
        "confidence": round(
            result["scores"][0],
            3
        )
    }



# -----------------------------
# 3. Entity Extraction
# -----------------------------

def extract_entities(query):

    doc = nlp(query)

    entities = []

    for ent in doc.ents:

        entities.append(
            {
                "text": ent.text,
                "type": ent.label_
            }
        )

    return entities



# -----------------------------
# 4. Complexity Detection
# -----------------------------

def detect_complexity(query):

    words=query.split()

    multi_words=[
        "compare",
        "difference",
        "and",
        "vs",
        "recommend"
    ]


    if len(words)>15:
        return "complex"


    for word in multi_words:
        if word in query.lower():
            return "multi-hop"


    return "simple"



# -----------------------------
# 5. Query Rewriting
# -----------------------------

def rewrite_query(query):

    prompt=f"""
Rewrite this query for better search retrieval.

Keep the meaning same.

Query:
{query}

Better query:
"""


    result=rewriter(
        prompt,
        max_new_tokens=50
    )


    return result[0]["generated_text"]



# -----------------------------
# 6. Query Expansion
# -----------------------------

def expand_query(query):

    related_words={
        "rag":[
            "retrieval augmented generation",
            "vector database",
            "document retrieval"
        ],

        "llm":[
            "large language model",
            "generative AI"
        ],

        "ai":[
            "machine learning",
            "deep learning"
        ]
    }


    expansion=[]

    for key,value in related_words.items():

        if key in query.lower():

            expansion.extend(value)


    return expansion



# -----------------------------
# 7. Query Decomposition
# -----------------------------

def decompose_query(query):

    keywords=[
        "compare",
        "and",
        "vs",
        "difference"
    ]


    for key in keywords:

        if key in query.lower():

            parts=query.split(key)

            return [
                part.strip()
                for part in parts
                if part.strip()
            ]


    return [query]



# -----------------------------
# 8. Query Routing
# -----------------------------
#
# Exactly 3 agents. classify_intent() already returns one of
# INTENT_LABELS, so this is a direct 1:1 map with a safe default
# rather than a many-to-one bucket of older intent names.

ROUTES = {
    "document lookup": "Document Agent",
    "database query": "SQL Agent",
    "web search": "Web Search Agent",
}

def route_query(intent):

    return ROUTES.get(
        intent,
        "Document Agent"   # default fallback if confidence is low / unclear
    )



# -----------------------------
# MAIN PIPELINE
# -----------------------------

def query_understanding(query):


    result={}


    # preprocessing
    clean_query=preprocess(query)


    result["original_query"]=query

    result["clean_query"]=clean_query


    # intent

    intent=classify_intent(
        clean_query
    )

    result["intent"]=intent



    # entities

    result["entities"]=extract_entities(
        query
    )



    # complexity

    result["complexity"]=detect_complexity(
        query
    )



    # rewriting

    result["rewritten_query"]=rewrite_query(
        query
    )



    # expansion

    result["expanded_terms"]=expand_query(
        query
    )



    # decomposition

    result["sub_queries"]=decompose_query(
        query
    )



    # routing

    result["route"]=route_query(
        intent["intent"]
    )


    return result



# -----------------------------
# TEST
# -----------------------------


if __name__ == "__main__":


        query = "What is Matma Gandi?"

        output = query_understanding(query)


        print("\n=========================================")
        print("QUERY:", query.strip())
        print("=========================================")

        for key, value in output.items():
            print("\n", key)
            print("----------------")
            print(value)