Modular RAG

NAIVE RAG (Monolithic):
Query → Embed → Search → LLM → Answer

MODULAR RAG (Flexible):
Query → [Query Rewriter] → [Router] → [Retriever A/B/C] 
      → [Reranker] → [Generator] → [Validator] → Answer

def rewrite_query(user_query: str, chat_history: list) -> str:
    prompt = f"""
    Dựa vào lịch sử chat, viết lại câu hỏi của user thành câu đầy đủ.
    
    History: {chat_history[-3:]}
    User: {user_query}
    
    Rewritten query:
    """
    return llm.generate(prompt)
 
# Example
history = ["User: Tôi muốn mua iPhone", "AI: Bạn muốn mua iPhone nào?"]
query = "Giá bao nhiêu?"
 
rewritten = rewrite_query(query, history)
# Output: "iPhone 15 Pro Max giá bao nhiêu tiền?"

from enum import Enum
 
class DataSource(Enum):
    PRODUCT_DB = "product_database"
    SUPPORT_WIKI = "support_wiki"
    CODE_REPO = "github_repo"
    WEB_SEARCH = "google_search"
    
def route_query(query: str) -> DataSource:
    prompt = f"""
    Phân loại query sau vào một trong các nguồn:
    - PRODUCT_DB: Câu hỏi về sản phẩm, giá, tồn kho
    - SUPPORT_WIKI: Câu hỏi về hỗ trợ, warranty, trả hàng
    - CODE_REPO: Câu hỏi về code, API, technical docs
    - WEB_SEARCH: Câu hỏi chung, tin tức, không trong database
    
    Query: "{query}"
    
    Trả về chỉ 1 từ: PRODUCT_DB, SUPPORT_WIKI, CODE_REPO, hoặc WEB_SEARCH
    """
    result = llm.generate(prompt).strip()
    return DataSource[result]
 
# Route to different retrievers
source = route_query("API rate limit là bao nhiêu?")
# Output: DataSource.CODE_REPO

class MultiRetriever:
    def __init__(self):
        self.retrievers = {
            DataSource.PRODUCT_DB: ProductDBRetriever(),
            DataSource.SUPPORT_WIKI: WikiRetriever(),
            DataSource.CODE_REPO: GithubRetriever(),
            DataSource.WEB_SEARCH: GoogleSearchRetriever(),
        }
    
    def retrieve(self, query: str, source: DataSource, k: int = 5):
        return self.retrievers[source].search(query, k)

def adaptive_retrieve(query: str, max_k: int = 10):
    results = vector_db.search(query, k=max_k)
    
    # Chỉ giữ docs có similarity > threshold
    threshold = 0.8
    filtered = [r for r in results if r.score > threshold]
    
    # Minimum 2, maximum 5
    return filtered[:5] if len(filtered) > 5 else (filtered if len(filtered) >= 2 else results[:2])

def select_generator(query: str, docs: list) -> str:
    """Chọn model dựa trên độ phức tạp."""
    
    # Tính complexity score
    complexity = estimate_complexity(query, docs)
    
    if complexity < 3:
        return "gpt-4o-mini"  # Nhanh, rẻ
    elif complexity < 7:
        return "gpt-4o"  # Balanced
    else:
        return "claude-3.5-sonnet"  # Best reasoning

def validate_response(query: str, response: str, docs: list) -> dict:
    prompt = f"""
    Đánh giá response sau:
    
    Query: {query}
    Response: {response}
    Source docs: {docs}
    
    Trả về JSON:
    {{
        "is_grounded": true/false,  // Response có dựa trên docs không?
        "is_relevant": true/false,  // Response có trả lời query không?
        "confidence": 0-100
    }}
    """
    return json.loads(llm.generate(prompt))

class ModularRAG:
    def __init__(self):
        self.query_rewriter = QueryRewriter()
        self.router = QueryRouter()
        self.multi_retriever = MultiRetriever()
        self.reranker = CohereReranker()
        self.generators = GeneratorPool()
        self.validator = ResponseValidator()
    
    def answer(self, query: str, history: list = None):
        # 1. Rewrite query with context
        enhanced_query = self.query_rewriter.rewrite(query, history)
        
        # 2. Route to data source
        source = self.router.route(enhanced_query)
        
        # 3. Retrieve from selected source
        docs = self.multi_retriever.retrieve(enhanced_query, source, k=20)
        
        # 4. Rerank
        top_docs = self.reranker.rerank(enhanced_query, docs, top_k=5)
        
        # 5. Select generator
        model = self.select_generator(enhanced_query, top_docs)
        
        # 6. Generate
        response = self.generators.generate(model, enhanced_query, top_docs)
        
        # 7. Validate
        validation = self.validator.validate(query, response, top_docs)
        
        if validation["confidence"] < 50:
            return self.fallback_response()
        
        return response

from openai import OpenAI
 
client = OpenAI()
 
CATEGORIES = {
    "product": "Câu hỏi về sản phẩm, giá cả, tính năng",
    "support": "Câu hỏi về hỗ trợ, bảo hành, trả hàng",
    "general": "Câu hỏi chung, chào hỏi, off-topic"
}
 
def route_query(query: str) -> str:
    prompt = f"""
    Phân loại query vào một trong các category:
    {CATEGORIES}
    
    Query: "{query}"
    
    Trả về chỉ 1 từ: product, support, hoặc general
    """
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=10
    )
    
    return response.choices[0].message.content.strip().lower()
 
# Test
test_queries = [
    "iPhone 15 Pro giá bao nhiêu?",
    "Làm sao đổi trả hàng?",
    "Chào bạn, hôm nay thế nào?"
]
 
for q in test_queries:
    category = route_query(q)
    print(f"'{q}' → {category}")