from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.response_synthesizers import get_response_synthesizer
from my_logging import log_message

SIMPLE_PROMPT = """Вы - эксперт по нормативной документации.

Контекст:
{context_str}

Вопрос: {query_str}

Инструкция:
1. Отвечайте ТОЛЬКО на основе предоставленного контекста
2. Цитируйте конкретные источники (документ, раздел, таблицу)
3. Если информации недостаточно, четко укажите это
4. Будьте точны и конкретны

Ответ:"""

def create_vector_index(documents):
    """Create vector index from documents"""
    log_message(f"Building vector index from {len(documents)} documents...")
    index = VectorStoreIndex.from_documents(documents)
    log_message("✓ Index created")
    return index

def create_query_engine(vector_index):
    """Create hybrid retrieval engine"""
    log_message("Creating query engine...")
    
    # Vector retriever
    vector_retriever = VectorIndexRetriever(
        index=vector_index,
        similarity_top_k=50
    )
    
    # BM25 retriever
    bm25_retriever = BM25Retriever.from_defaults(
        docstore=vector_index.docstore,
        similarity_top_k=50
    )
    
    # Hybrid fusion
    hybrid_retriever = QueryFusionRetriever(
        [vector_retriever, bm25_retriever],
        similarity_top_k=60,
        num_queries=1
    )
    
    # Response synthesizer
    response_synthesizer = get_response_synthesizer()
    
    # Query engine
    query_engine = RetrieverQueryEngine(
        retriever=hybrid_retriever,
        response_synthesizer=response_synthesizer
    )
    
    log_message("✓ Query engine created")
    return query_engine