index_retriever.py

from llama_index.core import VectorStoreIndex, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.llms.openai import OpenAI
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.response_synthesizers import get_response_synthesizer, ResponseMode
from llama_index.core.prompts import PromptTemplate
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever
from sentence_transformers import CrossEncoder
import logging

logger = logging.getLogger(__name__)

def log_message(message):
    logger.info(message)
    print(message, flush=True)

class IndexRetriever:
    def __init__(self, config):
        self.config = config
        self.vector_index = None
        self.query_engine = None
        self.reranker = None
        self.current_model = config.DEFAULT_MODEL
        
    def get_llm_model(self, model_name):
        try:
            model_config = self.config.AVAILABLE_MODELS.get(model_name)
            if not model_config:
                log_message(f"Модель {model_name} не найдена, использую модель по умолчанию")
                model_config = self.config.AVAILABLE_MODELS[self.config.DEFAULT_MODEL]
            
            if not model_config.get("api_key"):
                raise Exception(f"API ключ не найден для модели {model_name}")
            
            if model_config["provider"] == "google":
                return GoogleGenAI(
                    model=model_config["model_name"], 
                    api_key=model_config["api_key"]
                )
            elif model_config["provider"] == "openai":
                return OpenAI(
                    model=model_config["model_name"],
                    api_key=model_config["api_key"]
                )
            else:
                raise Exception(f"Неподдерживаемый провайдер: {model_config['provider']}")
                
        except Exception as e:
            log_message(f"Ошибка создания модели {model_name}: {str(e)}")
            return GoogleGenAI(model="gemini-2.0-flash", api_key=self.config.GOOGLE_API_KEY)

    def initialize_models(self, documents):
        try:
            log_message("Инициализация моделей и индекса")
            
            embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
            llm = self.get_llm_model(self.current_model)
            
            log_message("Инициализирую переранкер")
            self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
            
            Settings.embed_model = embed_model
            Settings.llm = llm
            
            log_message(f"Строю векторный индекс из {len(documents)} документов")
            self.vector_index = VectorStoreIndex.from_documents(documents)
            
            self.create_query_engine()
            
            log_message(f"Модели и индекс успешно инициализированы с моделью: {self.current_model}")
            return True
            
        except Exception as e:
            log_message(f"Ошибка инициализации моделей: {str(e)}")
            return False

    def create_query_engine(self):
        try:
            bm25_retriever = BM25Retriever.from_defaults(
                docstore=self.vector_index.docstore, 
                similarity_top_k=15
            )
            
            vector_retriever = VectorIndexRetriever(
                index=self.vector_index, 
                similarity_top_k=20,
                similarity_cutoff=0.5
            )
            
            hybrid_retriever = QueryFusionRetriever(
                [vector_retriever, bm25_retriever],
                similarity_top_k=30,
                num_queries=1
            )
            
            custom_prompt_template = PromptTemplate(self.config.CUSTOM_PROMPT)
            response_synthesizer = get_response_synthesizer(
                response_mode=ResponseMode.TREE_SUMMARIZE,
                text_qa_template=custom_prompt_template
            )
            
            self.query_engine = RetrieverQueryEngine(
                retriever=hybrid_retriever,
                response_synthesizer=response_synthesizer
            )
            
            log_message("Query engine успешно создан")
            
        except Exception as e:
            log_message(f"Ошибка создания query engine: {str(e)}")
            raise

    def switch_model(self, model_name):
        try:
            log_message(f"Переключение на модель: {model_name}")
            
            new_llm = self.get_llm_model(model_name)
            Settings.llm = new_llm
            
            if self.vector_index is not None:
                self.create_query_engine()
                self.current_model = model_name
                log_message(f"Модель успешно переключена на: {model_name}")
                return f"✅ Модель переключена на: {model_name}"
            else:
                return "❌ Ошибка: система не инициализирована"
                
        except Exception as e:
            error_msg = f"Ошибка переключения модели: {str(e)}"
            log_message(error_msg)
            return f"❌ {error_msg}"

    def rerank_nodes(self, query, nodes, top_k=10):
        if not nodes or not self.reranker:
            return nodes[:top_k]
        
        try:
            log_message(f"Переранжирую {len(nodes)} узлов")
            
            pairs = []
            for node in nodes:
                pairs.append([query, node.text])
            
            scores = self.reranker.predict(pairs)
            
            scored_nodes = list(zip(nodes, scores))
            scored_nodes.sort(key=lambda x: x[1], reverse=True)
            
            reranked_nodes = [node for node, score in scored_nodes[:top_k]]
            log_message(f"Возвращаю топ-{len(reranked_nodes)} переранжированных узлов")
            
            return reranked_nodes
        except Exception as e:
            log_message(f"Ошибка переранжировки: {str(e)}")
            return nodes[:top_k]

    def retrieve_nodes(self, question):
        if self.query_engine is None:
            return []
        
        try:
            log_message(f"Извлекаю релевантные узлы для вопроса: {question}")
            retrieved_nodes = self.query_engine.retriever.retrieve(question)
            log_message(f"Извлечено {len(retrieved_nodes)} узлов")
            
            log_message("Применяю переранжировку")
            reranked_nodes = self.rerank_nodes(question, retrieved_nodes, top_k=10)
            
            return reranked_nodes
            
        except Exception as e:
            log_message(f"Ошибка извлечения узлов: {str(e)}")
            return []

    def get_current_model(self):
        return self.current_model

    def is_initialized(self):
        return self.query_engine is not None