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 # 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', "AIzaSyDemsCp7JIdRNDRyP6DkYdMox1DLZwPcPE") 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 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