RAG_AIEXP_01 / index_retriever.py
MrSimple07's picture
simplest version
9985d37
Raw
History Blame
2.11 kB
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=30
)
# BM25 retriever
bm25_retriever = BM25Retriever.from_defaults(
docstore=vector_index.docstore,
similarity_top_k=30
)
# Hybrid fusion
hybrid_retriever = QueryFusionRetriever(
[vector_retriever, bm25_retriever],
similarity_top_k=40,
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