Multi-Query Rewriter RAG

# Template 1: Standalone, specific rewrite
REPHRASE_TEMPLATE_1 = """
Rewrite this question to be a standalone, specific query about pregnancy and childbirth.

Original question: {question}

Instructions:
- Maintain the medical/obstetric context if relevant.
- Be specific and clear in medical terms.
- Focus on pregnancy, childbirth, prenatal care, or maternal health.
- Ensure the question is complete and self-contained.

Standalone question:
"""

# Template 2: Synonym expansion
REPHRASE_TEMPLATE_2 = """
Rephrase this question about pregnancy and childbirth using synonyms and 
alternative medical terms.

Original question: {question}

Instructions:
- Use precise medical terminology.
- Include synonyms and alternative terms.
- Maintain the meaning but change the wording.
- Focus on clinical and obstetric aspects.

Rephrased question:
"""

# Template 3: Context expansion
REPHRASE_TEMPLATE_3 = """
Expand this question to include related aspects and additional context about 
pregnancy and childbirth.

Base question: {question}

Instructions:
- Expand the question to include related aspects.
- Add context about complications, prevention, or care.
- Include possible variations or special cases.
- Keep the focus on maternal and perinatal health.

Expanded question:
"""

def process_rewriter_query(
    question: str, 
    custom_rewriter_llm: ChatOpenAI = None, 
    custom_answer_llm: ChatOpenAI = None, 
    max_final_docs: int = 8
) -> Dict[str, Any]:
    """
    Processes a query using the multi-query rewriting RAG pipeline.
    
    Args:
        question (str): The user's question.
        custom_rewriter_llm (ChatOpenAI, optional): Custom model for query rewriting.
        custom_answer_llm (ChatOpenAI, optional): Custom model for answer generation.
        max_final_docs (int): The maximum number of documents to return.
    
    Returns:
        Dict[str, Any]: Answer, contexts, rewritten queries, and detailed metrics.
    """
    # 1. Generate rewritten queries and track metrics
    rewritten_queries = []
    rewrite_input_tokens, rewrite_output_tokens, rewrite_cost = 0, 0, 0
    
    for prompt in REPHRASE_PROMPTS:
        rewritten_query, rewrite_metrics = _invoke_text_with_usage(
            current_rewriter_llm,
            prompt.format(question=question)
        )
        rewritten_queries.append(rewritten_query)
        rewrite_input_tokens += rewrite_metrics["input_tokens"]
        rewrite_output_tokens += rewrite_metrics["output_tokens"]
        rewrite_cost += rewrite_metrics["cost"]
    
    # 2. Retrieve documents for each rewritten query
    all_docs_with_scores = []
    doc_ids_seen = set()
    
    for i, query in enumerate(rewritten_queries, 1):
        results = vectorstore.similarity_search_with_score(query, k=5)
        for doc, distance in results:
            similarity = max(0.0, 1.0 - distance)
            doc_id = doc.page_content[:100]  # Use content prefix as ID
            
            if doc_id not in doc_ids_seen:
                doc_ids_seen.add(doc_id)
                # Penalize queries from later, more speculative prompts
                query_weight = 1.0 - (i - 1) * 0.05
                all_docs_with_scores.append((doc, similarity * query_weight))
    
    # 3. Re-rank and select the best documents
    all_docs_with_scores.sort(key=lambda x: x[1], reverse=True)
    retrieved_docs = [doc for doc, _ in all_docs_with_scores[:max_final_docs]]
    
    # 4. Format context and generate final answer
    formatted_context = format_docs(retrieved_docs)
    answer, answer_metrics = _invoke_text_with_usage(
        current_answer_llm,
        qa_prompt.format_messages(context=formatted_context, question=question)
    )
    
    # 5. Consolidate and return all information
    return {
        'answer': answer,
        'contexts': [doc.page_content for doc in retrieved_docs],
        'retrieved_documents': retrieved_docs,
        'rewritten_queries': rewritten_queries,
        'metrics': {...}
    }

from src.rag.rewriter import query_for_evaluation

# Basic usage with default models
result = query_for_evaluation(
    question="¿Cuáles son los síntomas del parto prematuro?"
)

# With custom models for each stage
result = query_for_evaluation(
    question="¿Qué es la diabetes gestacional?",
    rewriter_model="gpt-3.5-turbo",  # Query rewriting
    answer_model="gpt-4o"             # Final answer generation
)

# With custom LLM instances
from langchain_openai import ChatOpenAI
rewriter_llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.3)
answer_llm = ChatOpenAI(model_name="gpt-4o", temperature=0)

result = query_for_evaluation(
    question="¿Qué cuidados necesito en el embarazo?",
    custom_rewriter_llm=rewriter_llm,
    custom_answer_llm=answer_llm
)

{
    "question": str,
    "answer": str,
    "contexts": List[str],           # Up to 8 contexts
    "source_documents": List,
    "metadata": {
        "num_contexts": 8,
        "retrieval_method": "multi_query_rewrite",
        "rewrite_count": 3,
        "llm_model": "gpt-4o",
        "rewriter_model": "gpt-3.5-turbo",
        "provider": "openai",
        "execution_time": 5.82,
        "input_tokens": 3124,          # Total across all LLM calls
        "output_tokens": 487,
        "total_cost": 0.004523,
        "usage_source": "provider"
    }
}

# Conservative (more faithful rewrites)
llm_rewriter = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.1)

# Balanced (default)
llm_rewriter = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.3)

# Creative (more diverse rewrites)
llm_rewriter = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.5)

# Concise context (faster, cheaper LLM calls)
result = process_rewriter_query(question, max_final_docs=5)

# Balanced (default)
result = process_rewriter_query(question, max_final_docs=8)

# Comprehensive context (maximum information)
result = process_rewriter_query(question, max_final_docs=12)

# Current: Linear decay
query_weight = 1.0 - (i - 1) * 0.05  # 1.0, 0.95, 0.90

# Alternative: Exponential decay (more aggressive)
query_weight = 0.9 ** (i - 1)  # 1.0, 0.9, 0.81

# Alternative: Equal weighting (trust all rewrites equally)
query_weight = 1.0  # 1.0, 1.0, 1.0

# Add a fourth rewrite strategy focused on complications
COMPLICATIONS_TEMPLATE = """
Rewrite this pregnancy/childbirth question to focus on potential 
complications, risk factors, and warning signs.

Original question: {question}

Complication-focused question:
"""

REPHRASE_PROMPTS = [
    PromptTemplate.from_template(REPHRASE_TEMPLATE_1),
    PromptTemplate.from_template(REPHRASE_TEMPLATE_2),
    PromptTemplate.from_template(REPHRASE_TEMPLATE_3),
    PromptTemplate.from_template(COMPLICATIONS_TEMPLATE)  # Add custom
]

result['metrics'] = {
    'rewrite_input_tokens': 234,
    'rewrite_output_tokens': 156,
    'rewrite_cost': 0.000456,
    'answer_input_tokens': 2890,
    'answer_output_tokens': 331,
    'answer_cost': 0.004067,
    'total_input_tokens': 3124,
    'total_output_tokens': 487,
    'total_cost': 0.004523
}