Query Enhancement

1from langchain_openai import ChatOpenAI, OpenAIEmbeddings
2from langchain_core.prompts import ChatPromptTemplate
3from langchain_core.output_parsers import StrOutputParser
4
5llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
6embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
7
8# Step 1: Generate hypothetical document
9hyde_prompt = ChatPromptTemplate.from_template(
10    """Given the question below, write a detailed paragraph that would
11    answer this question. Write as if it's from an official document.
12    
13    Question: {question}
14    
15    Hypothetical answer:"""
16)
17
18hyde_chain = hyde_prompt | llm | StrOutputParser()
19
20# Step 2: Use hypothetical answer for search
21def hyde_search(query, vectorstore, k=5):
22    # Generate hypothetical answer
23    hypothetical_doc = hyde_chain.invoke({"question": query})
24    print(f"HyDE document: {hypothetical_doc[:200]}...")
25    
26    # Search with hypothetical document embedding
27    results = vectorstore.similarity_search(hypothetical_doc, k=k)
28    return results
29
30# Usage
31results = hyde_search("lương tối thiểu 2024", vectorstore)

1from langchain.retrievers import MultiQueryRetriever
2
3# Auto-generate query variants
4multiquery_retriever = MultiQueryRetriever.from_llm(
5    retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
6    llm=ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
7)
8
9# Search with multiple queries automatically
10results = multiquery_retriever.invoke("RAG performance optimization")
11print(f"Retrieved {len(results)} unique documents")

1from langchain_core.prompts import ChatPromptTemplate
2
3multi_query_prompt = ChatPromptTemplate.from_template(
4    """You are an AI language model assistant. Your task is to generate 3
5    different versions of the given user question to retrieve relevant documents.
6    
7    Provide these alternative questions separated by newlines.
8    
9    Original question: {question}
10    
11    Alternative questions:"""
12)
13
14def custom_multi_query(query, vectorstore, k=3):
15    # Generate alternative queries
16    response = (multi_query_prompt | llm | StrOutputParser()).invoke({"question": query})
17    alt_queries = [q.strip() for q in response.strip().split("\n") if q.strip()]
18    
19    all_queries = [query] + alt_queries
20    print(f"Searching with {len(all_queries)} queries:")
21    for q in all_queries:
22        print(f"  - {q}")
23    
24    # Search each query
25    all_docs = []
26    seen_contents = set()
27    for q in all_queries:
28        docs = vectorstore.similarity_search(q, k=k)
29        for doc in docs:
30            content_hash = hash(doc.page_content)
31            if content_hash not in seen_contents:
32                seen_contents.add(content_hash)
33                all_docs.append(doc)
34    
35    print(f"Total unique documents: {len(all_docs)}")
36    return all_docs

1stepback_prompt = ChatPromptTemplate.from_template(
2    """You are an expert at generating step-back questions.
3    Given a specific question, generate a more general question that
4    would help retrieve background information.
5    
6    Original: {question}
7    Step-back question:"""
8)
9
10def stepback_search(query, vectorstore, k=3):
11    # Generate step-back question
12    stepback_q = (stepback_prompt | llm | StrOutputParser()).invoke({"question": query})
13    print(f"Step-back: {stepback_q}")
14    
15    # Search with both
16    original_docs = vectorstore.similarity_search(query, k=k)
17    stepback_docs = vectorstore.similarity_search(stepback_q, k=k)
18    
19    # Combine and deduplicate
20    all_docs = original_docs + stepback_docs
21    unique_docs = list({doc.page_content: doc for doc in all_docs}.values())
22    
23    return unique_docs

1decomposition_prompt = ChatPromptTemplate.from_template(
2    """Break down this complex question into 2-4 simpler sub-questions
3    that can be answered independently.
4    
5    Complex question: {question}
6    
7    Sub-questions (one per line):"""
8)
9
10def decompose_and_search(query, vectorstore, k=3):
11    # Decompose
12    response = (decomposition_prompt | llm | StrOutputParser()).invoke({"question": query})
13    sub_queries = [q.strip().lstrip("0123456789.-) ") for q in response.strip().split("\n") if q.strip()]
14    
15    print(f"Decomposed into {len(sub_queries)} sub-queries:")
16    
17    all_context = []
18    for i, sub_q in enumerate(sub_queries):
19        print(f"  {i+1}. {sub_q}")
20        docs = vectorstore.similarity_search(sub_q, k=k)
21        context = "\n".join([d.page_content for d in docs])
22        all_context.append(f"Sub-question: {sub_q}\nContext: {context}")
23    
24    return "\n\n---\n\n".join(all_context)
25
26# Usage
27complex_query = "So sánh lương tối thiểu vùng 1 và vùng 4, và cách tính bảo hiểm dựa trên lương"
28context = decompose_and_search(complex_query, vectorstore)

1from langchain_core.prompts import ChatPromptTemplate
2from langchain_core.output_parsers import StrOutputParser
3
4router_prompt = ChatPromptTemplate.from_template(
5    """Given the user question, classify which knowledge base to search.
6    
7    Available knowledge bases:
8    - "hr_policy": HR policies, leave, benefits, workplace rules
9    - "technical": Technical documentation, APIs, system guides
10    - "legal": Legal regulations, compliance, contracts
11    - "general": General company information
12    
13    Question: {question}
14    
15    Answer with just the knowledge base name:"""
16)
17
18def route_query(query, vectorstores):
19    """Route query to appropriate vector store."""
20    kb_name = (router_prompt | llm | StrOutputParser()).invoke({"question": query})
21    kb_name = kb_name.strip().lower()
22    
23    if kb_name in vectorstores:
24        print(f"Routing to: {kb_name}")
25        return vectorstores[kb_name].similarity_search(query, k=5)
26    else:
27        print(f"Unknown KB '{kb_name}', searching all")
28        all_results = []
29        for vs in vectorstores.values():
30            all_results.extend(vs.similarity_search(query, k=2))
31        return all_results
32
33# Usage
34vectorstores = {
35    "hr_policy": hr_vectorstore,
36    "technical": tech_vectorstore,
37    "legal": legal_vectorstore
38}
39docs = route_query("Chính sách nghỉ phép năm", vectorstores)

1class EnhancedRetriever:
2    """Combine multiple query enhancement strategies."""
3    
4    def __init__(self, vectorstore, llm):
5        self.vectorstore = vectorstore
6        self.llm = llm
7    
8    def retrieve(self, query, strategy="auto", k=5):
9        if strategy == "auto":
10            strategy = self._detect_strategy(query)
11        
12        print(f"Using strategy: {strategy}")
13        
14        if strategy == "hyde":
15            return self._hyde_search(query, k)
16        elif strategy == "multi_query":
17            return self._multi_query_search(query, k)
18        elif strategy == "stepback":
19            return self._stepback_search(query, k)
20        elif strategy == "decompose":
21            return self._decompose_search(query, k)
22        else:
23            return self.vectorstore.similarity_search(query, k=k)
24    
25    def _detect_strategy(self, query):
26        """Auto-detect best strategy based on query."""
27        word_count = len(query.split())
28        
29        if " và " in query or " so sánh " in query.lower():
30            return "decompose"
31        elif word_count <= 3:
32            return "hyde"
33        elif "?" in query:
34            return "multi_query"
35        else:
36            return "stepback"
37
38# Usage
39retriever = EnhancedRetriever(vectorstore, llm)
40docs = retriever.retrieve("lương tối thiểu", strategy="auto")