scripts/config.py

import os
import google.generativeai as genai
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core import Settings
from llama_index.core.llms import ChatMessage, MessageRole
import os
from huggingface_hub import hf_hub_download
from huggingface_hub import HfApi, login, create_repo



# Configuration
EMBEDDING_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
RETRIEVER_TOP_K = 10
RETRIEVER_SIMILARITY_CUTOFF = 0.7
RAG_FILES_DIR = "processed_data"
PROCESSED_DATA_FILE = "processed_data/processed_chunks.csv"

UPLOAD_FOLDER = "UPLOADED_DOCUMENTS"
INDEX_STATE_FILE = "processed_data/index_store.json"

GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
HF_REPO_ID = "MrSimple01/AIEXP_RAG_FILES"
HF_TOKEN = os.getenv('HF_TOKEN')
LLM_MODEL = "gemini-2.0-flash"

CHUNK_SIZE = 1024
CHUNK_OVERLAP = 256
MAX_CHUNK_SIZE = 2048
MIN_CHUNK_SIZE = 750
SIMILARITY_THRESHOLD = 0.7


def login_to_huggingface():
    """Login to HuggingFace Hub"""
    try:
        if HF_TOKEN:
            login(token=HF_TOKEN)
            print("✅ Successfully logged in to HuggingFace Hub with token")
            return True
        else:
            # Interactive login if no token provided
            login()
            print("✅ Successfully logged in to HuggingFace Hub")
            return True
    except Exception as e:
        print(f"❌ Failed to login to HuggingFace Hub: {e}")
        return False


def download_pretrained_files():
    """Download pre-trained RAG files from HuggingFace Hub"""
    try:
        print("Downloading pre-trained RAG files from HuggingFace Hub...")
        
        # Files to download
        files_to_download = [
            "faiss_index.index",
            "processed_chunks.csv",
            "chunk_metadata.pkl",
            "config.pkl",
            "documents.pkl",
            "default__vector_store.json",
            "docstore.json",
            "index_store.json"
        ]
        
        # Ensure RAG_FILES_DIR exists
        os.makedirs(RAG_FILES_DIR, exist_ok=True)
        os.makedirs("processed_data", exist_ok=True)
        
        downloaded_files = {}
        
        for filename in files_to_download:
            try:
                print(f"Downloading {filename}...")
                
                # Download to RAG_FILES_DIR for most files, processed_data for CSV
                target_dir = "processed_data" if filename == "processed_chunks.csv" else RAG_FILES_DIR
                
                file_path = hf_hub_download(
                    repo_id=HF_REPO_ID,
                    filename=filename,
                    local_dir=target_dir,
                    repo_type="dataset",
                    token=HF_TOKEN
                )
                
                downloaded_files[filename] = file_path
                print(f"✓ Downloaded {filename}")
                
            except Exception as e:
                print(f"✗ Failed to download {filename}: {e}")
                continue
        
        # Verify critical files
        critical_files = ["faiss_index.index", "processed_chunks.csv"]
        missing_critical = [f for f in critical_files if f not in downloaded_files]
        
        if missing_critical:
            print(f"❌ Missing critical files: {missing_critical}")
            return False
        
        print(f"✅ Successfully downloaded {len(downloaded_files)}/{len(files_to_download)} files")
        return True
        
    except Exception as e:
        print(f"❌ Failed to download pre-trained files: {e}")
        return False

def setup_llm_settings():
    """Setup embedding and LLM models"""
    # Configure Google API
    if GOOGLE_API_KEY:
        genai.configure(api_key=GOOGLE_API_KEY)
    
    # Set embedding model
    embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL)
    Settings.embed_model = embed_model
    
    # Set LLM - IMPORTANT: This prevents OpenAI default
    if GOOGLE_API_KEY:
        try:
            llm = GoogleGenAI(model=LLM_MODEL, api_key=GOOGLE_API_KEY)
            Settings.llm = llm
            print("Google GenAI LLM initialized successfully")
        except Exception as e:
            print(f"Warning: Could not initialize Google GenAI LLM: {e}")
            # Set a dummy LLM to prevent OpenAI default
            from llama_index.core.llms.mock import MockLLM
            Settings.llm = MockLLM()
    else:
        print("Warning: GOOGLE_API_KEY not found. Using MockLLM.")
        from llama_index.core.llms.mock import MockLLM
        Settings.llm = MockLLM()



CUSTOM_PROMPT = """
Вы являетесь высокоспециализированным Ассистентом для анализа документов (AIEXP). Ваша цель - предоставлять точные, корректные и контекстно релевантные ответы на основе анализа нормативной документации (НД). Все ваши ответы должны основываться исключительно на предоставленном контексте без использования внешних знаний или предположений.

ОПРЕДЕЛЕНИЕ ТИПА ЗАДАЧИ:
Проанализируйте запрос пользователя и определите тип задачи:

1. КРАТКОЕ САММАРИ (ключевые слова: "кратко", "суммировать", "резюме", "основные моменты", "в двух словах"):
   - Предоставьте структурированное резюме запрашиваемого раздела/пункта
   - Выделите ключевые требования, процедуры или положения
   - Используйте нумерованный список для лучшей читаемости
   - Сохраняйте терминологию НД

2. ПОИСК ДОКУМЕНТА И ПУНКТА (ключевые слова: "найти", "где", "какой документ", "в каком разделе", "ссылка"):
   - Укажите конкретный документ и его структурное расположение
   - Предоставьте точные номера разделов/подразделов/пунктов
   - Процитируйте релевантные фрагменты
   - Если найдено несколько документов, перечислите все с указанием специфики каждого

3. ПРОВЕРКА КОРРЕКТНОСТИ (ключевые слова: "правильно ли", "соответствует ли", "проверить", "корректно", "нарушение"):
   - Сопоставьте предоставленную информацию с требованиями НД
   - Четко укажите: "СООТВЕТСТВУЕТ" или "НЕ СООТВЕТСТВУЕТ"
   - Перечислите конкретные требования НД
   - Укажите выявленные расхождения или подтвердите соответствие
   - Процитируйте релевантные пункты НД

4. ПЛАН ДЕЙСТВИЙ (ключевые слова: "план", "алгоритм", "последовательность", "как действовать", "пошагово"):
   - Создайте пронумерованный пошаговый план
   - Каждый шаг должен содержать ссылку на соответствующий пункт НД
   - Укажите необходимые документы или формы
   - Добавьте временные рамки, если они указаны в НД
   - Выделите критические требования или ограничения

ПРАВИЛА ФОРМИРОВАНИЯ ОТВЕТОВ:

1. ОБЯЗАТЕЛЬНОЕ УКАЗАНИЕ ИСТОЧНИКОВ:
   - Для контента из конкретного раздела/подраздела:
     "Согласно разделу [X] и подразделу [X.X]: [Ваш ответ]"
   - Для контента вне подразделов (таблицы, рисунки, общие разделы):
     "Согласно [Название документа] - [Номер и наименование пункта/таблицы/рисунка]: [Ваш ответ]"
   - При наличии метаданных о разделе и подразделе - включайте оба
   - При наличии только раздела: "Согласно разделу [X]: [Ваш ответ]"

2. СТРОГОЕ СЛЕДОВАНИЕ КОНТЕКСТУ:
   - Если информация не найдена: "Информация по вашему запросу не была найдена в нормативной документации."
   - Не делайте предположений или выводов за пределами предоставленного контекста
   - Не используйте общие знания

3. ИСПОЛЬЗОВАНИЕ ТЕРМИНОЛОГИИ НД:
   - Применяйте официальную терминологию из документов
   - Сохраняйте оригинальные формулировки ключевых требований
   - При необходимости разъясняйте специальные термины на основе НД

4. СТРУКТУРИРОВАНИЕ ОТВЕТОВ:
   - Для саммари: используйте маркированные или нумерованные списки
   - Для проверки: четкая структура "Требование → Соответствие/Несоответствие"
   - Для планов: пронумерованные шаги с подзадачами при необходимости
   - Для поиска: указание иерархии документа

5. ДОПОЛНИТЕЛЬНЫЕ РЕКОМЕНДАЦИИ:
   - При множественных релевантных источниках - укажите все
   - Выделяйте критически важные требования
   - Указывайте альтернативные процедуры, если они предусмотрены НД

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

Вопрос: {query_str}

Ответ:
"""


LLM_MODEL_PREPROCESS = "gemini-1.5-flash"

def preprocess_query_with_context(user_query, chat_history=None, llm=None):
    if not chat_history:
        return user_query
    
    if not llm:
        llm = GoogleGenAI(model=LLM_MODEL_PREPROCESS, temperature=0.1)

    # Format chat history into a string for the prompt
    history_context = "\n".join([
        f"User: {item['user']}\nAssistant: {item['assistant']}"
        for item in chat_history[-3:]  # Consider only the last 3 exchanges for conciseness
    ])

    preprocessing_prompt = f"""Analyze the user's current question in the context of their chat history and improve it for better document retrieval.

Chat History:
{history_context}

Current Question: {user_query}

Tasks:
1. If the question refers to previous context, make it self-contained.
2. Add relevant keywords that would help find documents.
3. Maintain the legal/regulatory focus.
4. Keep it concise but specific.

Return ONLY the improved question:
"""

    try:
        messages = [ChatMessage(role=MessageRole.USER, content=preprocessing_prompt)]
        response = llm.chat(messages)
        improved_query = response.message.content.strip()

        # Fallback to the original query if the preprocessing fails or provides an overly long response
        if len(improved_query) > len(user_query) * 3 or not improved_query:
            return user_query
        
        return improved_query
    except Exception as e:
        print(f"Query preprocessing failed: {e}")
        return user_query


def create_chat_context_prompt(base_response, chat_history=None):
    if not chat_history:
        return base_response

    base_aware_response = base_response

    if len(chat_history) > 0:
        last_exchange = chat_history[-1]
        if any(keyword in last_exchange['user'].lower() for keyword in ['закон', 'кодекс', 'статья']):
            # Add a conversational prefix
            base_aware_response = f"Продолжая тему нормативных документов: {base_response}"

    return base_aware_response