Spaces:
Sleeping
Sleeping
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 | |
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" | |
PROCESSED_DATA_FILE = "processed_data/processed_chunks.csv" | |
INDEX_STATE_FILE = "processed_data/index_store.json" | |
RAG_FILES_DIR = "processed_data" | |
GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY', "AIzaSyDemsCp7JIdRNDRyP6DkYdMox1DLZwPcPE") | |
LLM_MODEL = "gemini-2.0-flash" | |
CHUNK_SIZE = 1024 | |
CHUNK_OVERLAP = 256 | |
MAX_CHUNK_SIZE = 2048 | |
MIN_CHUNK_SIZE = 750 | |
SIMILARITY_THRESHOLD = 0.7 | |
RETRIEVER_TOP_K = 15 | |
RETRIEVER_SIMILARITY_CUTOFF = 0.7 | |
def setup_llm_settings(): | |
# Set embedding model first | |
Settings.embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL) | |
# Only set LLM if API key is available | |
if GOOGLE_API_KEY: | |
try: | |
llm = GoogleGenAI(model=LLM_MODEL, api_key=GOOGLE_API_KEY) | |
Settings.llm = llm | |
except Exception as e: | |
print(f"Warning: Could not initialize Google GenAI LLM: {e}") | |
Settings.llm = None | |
else: | |
print("Warning: GOOGLE_API_KEY not found. Setting LLM to None.") | |
Settings.llm = None | |
# Intent Classification Prompt for preprocessing | |
INTENT_CLASSIFIER_PROMPT = """Проанализируй запрос пользователя и определи, к какой из следующих категорий он относится: | |
1. SUMMARY - Запрос на краткое изложение, саммари пункта, раздела или документа | |
Примеры: "кратко расскажи о разделе 5", "саммари главы про безопасность", "что говорится в пункте 3.2" | |
2. SEARCH - Поиск конкретной информации, документа или пункта | |
Примеры: "где найти информацию о...", "какой документ регулирует...", "в каком пункте говорится про..." | |
3. VALIDATION - Проверка корректности данных, соответствия требованиям | |
Примеры: "правильно ли...", "соответствует ли требованиям...", "проверь корректность данных..." | |
4. ACTION_PLAN - Составление плана действий на основе документации | |
Примеры: "составь план...", "какие шаги нужно предпринять...", "алгоритм действий для..." | |
Запрос пользователя: {user_query} | |
Ответь ТОЛЬКО одним словом: SUMMARY, SEARCH, VALIDATION или ACTION_PLAN""" | |
# Task-specific prompts | |
SUMMARY_PROMPT = """Ты специализированный помощник по анализу нормативной документации (НД) для создания кратких изложений и саммари. | |
Твоя задача: Предоставить краткое, структурированное изложение запрашиваемой информации из нормативных документов. | |
Правила создания саммари: | |
1. Обязательное указание источника: Всегда указывай источник информации в формате: | |
- Для конкретного раздела/подраздела: Согласно разделу [X] и подразделу [X.X]: [Саммари] | |
- Для общих пунктов: Согласно [Название документа] - [Номер и наименование пункта]: [Саммари] | |
2. Структурированность: Организуй информацию логично: | |
- Основные положения | |
- Ключевые требования | |
- Важные исключения или особенности | |
3. Краткость и полнота: Включи все существенные моменты, но избегай избыточных деталей | |
4. Сохранение терминологии НД: Используй точную терминологию из документов | |
5. Если информация отсутствует: Информация для создания саммари по вашему запросу не найдена в нормативной документации. | |
Контекст: {context_str} | |
Запрос: {query_str} | |
Саммари:""" | |
SEARCH_PROMPT = """Ты специализированный помощник по поиску информации в нормативной документации (НД). | |
Твоя задача: Найти и предоставить точную информацию, документы или пункты согласно запросу пользователя. | |
Правила поиска информации: | |
1. Точное указание источника: Обязательно укажи точное местоположение информации: | |
- Для конкретного раздела/подраздела: Согласно разделу [X] и подразделу [X.X]: [Найденная информация] | |
- Для общих пунктов: Согласно [Название документа] - [Номер и наименование пункта]: [Найденная информация] | |
2. Релевантность: Предоставляй именно ту информацию, которая отвечает на запрос | |
3. Полнота ответа: Если найдено несколько релевантных источников, укажи все | |
4. Точность цитирования: Используй формулировки непосредственно из НД | |
5. Если информация не найдена: Информация по вашему запросу не была найдена в нормативной документации. | |
Контекст: {context_str} | |
Поисковый запрос: {query_str} | |
Результат поиска:""" | |
VALIDATION_PROMPT = """Ты специализированный помощник по проверке соответствия требованиям нормативной документации (НД). | |
Твоя задача: Проанализировать предоставленные данные или информацию на соответствие требованиям НД и дать заключение о корректности. | |
Правила проведения проверки: | |
1. Источник требований: Четко укажи, согласно каким пунктам НД проводится проверка: | |
- Согласно разделу [X] и подразделу [X.X]: [Требование] | |
- Согласно [Название документа] - [Номер и наименование пункта]: [Требование] | |
2. Структура заключения: | |
- СООТВЕТСТВУЕТ/НЕ СООТВЕТСТВУЕТ требованиям | |
- Обоснование заключения со ссылками на конкретные пункты НД | |
- При несоответствии - указание на конкретные нарушения | |
- Рекомендации по устранению несоответствий (если применимо) | |
3. Объективность: Основывайся исключительно на требованиях НД, без субъективных оценок | |
4. Если требования не найдены: Требования для проверки данной информации не найдены в нормативной документации. | |
Контекст: {context_str} | |
Запрос на проверку: {query_str} | |
Заключение о соответствии:""" | |
ACTION_PLAN_PROMPT = """Ты специализированный помощник по составлению планов действий на основе нормативной документации (НД). | |
Твоя задача: Создать пошаговый план действий, основанный на требованиях и процедурах, описанных в НД. | |
Правила составления плана действий: | |
1. Источник каждого шага: Для каждого шага указывай источник в НД: | |
- Согласно разделу [X] и подразделу [X.X]: Шаг [N] | |
- Согласно [Название документа] - [Номер и наименование пункта]: Шаг [N] | |
2. Структура плана: | |
- Нумерованные шаги в логической последовательности | |
- Четкие действия для каждого шага | |
- Указание ответственных лиц/органов (если указано в НД) | |
- Временные рамки (если указаны в НД) | |
- Необходимые документы/материалы | |
3. Практичность: План должен быть выполнимым и соответствовать реальным процедурам | |
4. Полнота: Включи все обязательные этапы согласно НД | |
5. Если информация недостаточна: Недостаточно информации в нормативной документации для составления полного плана действий по вашему запросу. | |
Контекст: {context_str} | |
Запрос на план: {query_str} | |
План действий:""" | |
# Dictionary mapping intents to prompts | |
TASK_PROMPTS = { | |
"SUMMARY": SUMMARY_PROMPT, | |
"SEARCH": SEARCH_PROMPT, | |
"VALIDATION": VALIDATION_PROMPT, | |
"ACTION_PLAN": ACTION_PLAN_PROMPT | |
} | |
LLM_MODEL_PREPROCESS = "gemini-1.5-flash" | |
def classify_intent(user_query, llm=None): | |
"""Classify user intent to determine appropriate prompt""" | |
if not llm: | |
llm = GoogleGenAI(model=LLM_MODEL_PREPROCESS, temperature=0.1) | |
try: | |
intent_prompt = INTENT_CLASSIFIER_PROMPT.format(user_query=user_query) | |
messages = [ChatMessage(role=MessageRole.USER, content=intent_prompt)] | |
response = llm.chat(messages) | |
intent = response.message.content.strip() | |
# Validate intent and return default if invalid | |
if intent in TASK_PROMPTS: | |
return intent | |
else: | |
print(f"Unknown intent: {intent}, defaulting to SEARCH") | |
return "SEARCH" | |
except Exception as e: | |
print(f"Intent classification failed: {e}, defaulting to SEARCH") | |
return "SEARCH" | |
def get_task_specific_prompt(intent, context_str, query_str): | |
"""Get the appropriate prompt template based on intent""" | |
prompt_template = TASK_PROMPTS.get(intent, SEARCH_PROMPT) | |
return prompt_template.format(context_str=context_str, query_str=query_str) | |
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 | |
# Example usage function | |
def process_user_query(user_query, context_str, chat_history=None): | |
# Step 1: Classify intent | |
intent = classify_intent(user_query) | |
print(f"Detected intent: {intent}") | |
# Step 2: Preprocess query if needed | |
processed_query = preprocess_query_with_context(user_query, chat_history) | |
# Step 3: Get appropriate prompt | |
task_prompt = get_task_specific_prompt(intent, context_str, processed_query) | |
return task_prompt, intent |