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)