Spaces:
Running
Running
import logging | |
import os | |
import json | |
import traceback | |
from pathlib import Path | |
import pandas as pd | |
import tiktoken | |
from typing import List, Dict, Any, Optional, Tuple | |
from llama_index.core import ( | |
Document, | |
) | |
# Встановлюємо змінну середовища, щоб примусово використовувати CPU | |
os.environ["CUDA_VISIBLE_DEVICES"] = "" | |
os.environ["TORCH_DEVICE"] = "cpu" | |
from modules.config.ai_settings import ( | |
get_metadata_csv, | |
CHUNK_SIZE, | |
CHUNK_OVERLAP, | |
EXCLUDED_EMBED_METADATA_KEYS, | |
EXCLUDED_LLM_METADATA_KEYS, | |
GOOGLE_EMBEDDING_MODEL | |
) | |
# Налаштування логування | |
logger = logging.getLogger(__name__) | |
def initialize_embedding_model(): | |
""" | |
Ініціалізує модель ембедингів згідно налаштувань. | |
Використовує офіційний пакет GeminiEmbedding для Google Embeddings. | |
Returns: | |
object: Модель ембедингів | |
""" | |
try: | |
# ПЕРША СПРОБА: Google Embeddings через офіційний пакет | |
google_api_key = os.getenv("GEMINI_API_KEY") | |
# Перевіряємо наявність API ключа | |
if google_api_key: | |
try: | |
logger.info("Спроба ініціалізації Google Embeddings API через GeminiEmbedding...") | |
from llama_index.embeddings.gemini import GeminiEmbedding | |
# Використовуємо модель Gemini для ембедингів | |
model_name = "models/embedding-004" # або "models/text-embedding-001" | |
# Створюємо модель ембедингів Gemini | |
embed_model = GeminiEmbedding( | |
model_name=model_name, | |
api_key=google_api_key, | |
task_type="retrieval_query" # або "retrieval_document" | |
) | |
# Тестуємо модель | |
logger.info("Виконуємо тестовий запит до Gemini Embeddings API...") | |
test_embedding = embed_model.get_text_embedding("Тестовий запит до Gemini Embeddings API") | |
if test_embedding: | |
logger.info(f"Тестовий запит успішний, отримано ембединг розмірністю {len(test_embedding)}") | |
logger.info(f"Успішно ініціалізовано модель ембедингів Google Gemini: {model_name}") | |
return embed_model | |
else: | |
raise Exception("Тестове підключення до Google API не вдалося - отримано порожній результат") | |
except ImportError as imp_err: | |
logger.error(f"Помилка імпорту модуля GeminiEmbedding: {imp_err}") | |
logger.error("Можливо, потрібно встановити пакет: pip install llama-index-embeddings-gemini") | |
logger.warning("Спробуємо альтернативні методи...") | |
# Спробуємо альтернативний імпорт для Google AI SDK | |
try: | |
# Через Google GenAI SDK безпосередньо | |
from google import genai | |
logger.info("Спроба ініціалізації через Google GenAI API безпосередньо...") | |
# Ініціалізуємо клієнт Google GenAI | |
genai.configure(api_key=google_api_key) | |
client = genai.Client() | |
# Функція для отримання ембедингів від Google API | |
def get_google_embeddings(texts): | |
if not isinstance(texts, list): | |
texts = [texts] | |
try: | |
# Використовуємо Google Embeddings API | |
result = client.models.embed_content( | |
model=GOOGLE_EMBEDDING_MODEL, | |
contents=texts, | |
config={"task_type": "retrieval_query"} | |
) | |
# Виймаємо ембединги | |
embeddings = [embedding.values for embedding in result.embeddings] | |
# Повертаємо в правильному форматі для LlamaIndex | |
return embeddings[0] if len(embeddings) == 1 else embeddings | |
except Exception as e: | |
logger.error(f"Помилка при отриманні ембедингів від Google API: {e}") | |
logger.error(traceback.format_exc()) | |
raise | |
# Тестуємо | |
test_result = get_google_embeddings(["Тестовий запит до Google GenAI API"]) | |
if test_result: | |
# Створюємо кастомну модель ембедингів | |
embed_model = CustomEmbedding( | |
embed_func=get_google_embeddings, | |
embed_batch_size=8 | |
) | |
logger.info(f"Успішно ініціалізовано кастомну модель ембедингів Google через GenAI SDK") | |
return embed_model | |
else: | |
raise Exception("Тестове підключення до Google API не вдалося") | |
except ImportError: | |
logger.error("Не вдалося імпортувати ні llama-index-embeddings-gemini, ні google.genai") | |
except Exception as e: | |
logger.error(f"Помилка при альтернативній ініціалізації Google Embeddings: {e}") | |
logger.error(traceback.format_exc()) | |
except Exception as e: | |
logger.error(f"Не вдалося ініціалізувати Google Embeddings API: {e}") | |
logger.error(traceback.format_exc()) | |
else: | |
logger.warning("API ключ Google не знайдено (змінна GOOGLE_API_KEY не встановлена)") | |
logger.warning("Будь ласка, додайте GOOGLE_API_KEY у файл .env або змінні середовища") | |
# ДРУГА СПРОБА: HuggingFace ембединги | |
logger.info("Використання локальних HuggingFace ембедингів...") | |
from llama_index.embeddings.huggingface import HuggingFaceEmbedding | |
from modules.config.ai_settings import DEFAULT_EMBEDDING_MODEL, FALLBACK_EMBEDDING_MODEL | |
try: | |
# Явно вказуємо використання CPU | |
embed_model = HuggingFaceEmbedding( | |
model_name=DEFAULT_EMBEDDING_MODEL, | |
device="cpu" # Явно вказуємо CPU | |
) | |
logger.info(f"Успішно ініціалізовано модель ембедингів HuggingFace на CPU: {DEFAULT_EMBEDDING_MODEL}") | |
return embed_model | |
except Exception as e: | |
logger.warning(f"Не вдалося ініціалізувати основну модель HuggingFace ембедингів: {e}") | |
# Спробуємо резервну модель | |
try: | |
embed_model = HuggingFaceEmbedding( | |
model_name=FALLBACK_EMBEDDING_MODEL, | |
device="cpu" # Явно вказуємо CPU | |
) | |
logger.info(f"Успішно ініціалізовано резервну модель HuggingFace ембедингів на CPU: {FALLBACK_EMBEDDING_MODEL}") | |
return embed_model | |
except Exception as fallback_error: | |
logger.error(f"Не вдалося ініціалізувати резервну модель HuggingFace: {fallback_error}") | |
# Створення найпростішого фальшивого ембедера для аварійної ситуації | |
try: | |
from llama_index.embeddings.custom import CustomEmbedding | |
except ImportError: | |
# Для сумісності зі старими версіями бібліотеки | |
from llama_index.core.embeddings.custom import CustomEmbedding | |
import numpy as np | |
def fallback_embedding_func(texts): | |
if not isinstance(texts, list): | |
texts = [texts] | |
# Генеруємо фіктивні ембедінги (розмірність 768 - типова) | |
embeddings = [np.random.rand(768).tolist() for _ in texts] | |
return embeddings[0] if len(embeddings) == 1 else embeddings | |
logger.warning("Використовуємо аварійний фальшивий ембедер") | |
return CustomEmbedding(embed_func=fallback_embedding_func) | |
except Exception as e: | |
logger.error(f"Критична помилка при ініціалізації моделей ембедингів: {e}") | |
logger.error(traceback.format_exc()) | |
# Аварійний фальшивий ембедер | |
try: | |
from llama_index.embeddings.custom import CustomEmbedding | |
except ImportError: | |
# Для сумісності зі старими версіями бібліотеки | |
from llama_index.core.embeddings.custom import CustomEmbedding | |
import numpy as np | |
def emergency_embedding_func(texts): | |
if not isinstance(texts, list): | |
texts = [texts] | |
return [np.random.rand(768).tolist() for _ in texts] | |
logger.warning("Використовуємо аварійний фальшивий ембедер через критичну помилку") | |
return CustomEmbedding(embed_func=emergency_embedding_func) | |
def count_tokens(text, model="gpt-3.5-turbo"): | |
""" | |
Підраховує приблизну кількість токенів для тексту. | |
Args: | |
text (str): Текст для підрахунку токенів | |
model (str): Назва моделі для вибору енкодера | |
Returns: | |
int: Кількість токенів | |
""" | |
try: | |
encoding = tiktoken.encoding_for_model(model) | |
tokens = encoding.encode(text) | |
return len(tokens) | |
except Exception as e: | |
logger.warning(f"Не вдалося підрахувати токени через tiktoken: {e}") | |
# Якщо не можемо використати tiktoken, робимо просту оцінку | |
return len(text) // 3 # Приблизна оцінка | |
def convert_dataframe_to_documents(df: pd.DataFrame) -> List[Document]: | |
""" | |
Перетворює DataFrame з даними Jira в документи для індексування. | |
Args: | |
df (pd.DataFrame): DataFrame з даними Jira | |
Returns: | |
List[Document]: Список документів для індексування | |
""" | |
logger.info("Перетворення даних DataFrame в документи для LlamaIndex...") | |
jira_documents = [] | |
total_tokens = 0 | |
for idx, row in df.iterrows(): | |
# Основний текст - опис тікета | |
text = "" | |
if 'Description' in row and pd.notnull(row['Description']): | |
text = str(row['Description']) | |
# Додавання коментарів, якщо вони є | |
for col in df.columns: | |
if col.startswith('Comment') and pd.notnull(row[col]): | |
text += f"\n\nКоментар: {str(row[col])}" | |
# Метадані для документа | |
metadata = metadata = get_metadata_csv(row, idx) | |
# Додатково перевіряємо поле зв'язків, якщо воно є | |
if 'Outward issue link (Relates)' in row and pd.notnull(row['Outward issue link (Relates)']): | |
metadata["related_issues"] = row['Outward issue link (Relates)'] | |
# Додатково перевіряємо інші можливі поля зв'язків | |
for col in df.columns: | |
if col.startswith('Outward issue link') and col != 'Outward issue link (Relates)' and pd.notnull(row[col]): | |
link_type = col.replace('Outward issue link (', '').replace(')', '') | |
if "links" not in metadata: | |
metadata["links"] = {} | |
metadata["links"][link_type] = str(row[col]) | |
# Створюємо документ з вказаними виключеннями | |
doc = Document( | |
text=text, | |
metadata=metadata, | |
excluded_embed_metadata_keys=EXCLUDED_EMBED_METADATA_KEYS, | |
excluded_llm_metadata_keys=EXCLUDED_LLM_METADATA_KEYS | |
) | |
# Підраховуємо токени | |
token_count = count_tokens(text) | |
total_tokens += token_count | |
# Додаємо документ до списку | |
jira_documents.append(doc) | |
logger.info(f"Створено {len(jira_documents)} документів з {total_tokens} токенами") | |
return jira_documents | |
def check_index_integrity(indices_path: str) -> Tuple[bool, str]: | |
""" | |
Перевіряє цілісність індексів. | |
Args: | |
indices_path (str): Шлях до директорії з індексами | |
Returns: | |
Tuple[bool, str]: (True, '') якщо індекси валідні, (False, 'повідомлення про помилку') в іншому випадку | |
""" | |
try: | |
indices_path = Path(indices_path) | |
# Перевірка наявності директорії | |
if not indices_path.exists() or not indices_path.is_dir(): | |
return False, f"Директорія з індексами не існує: {indices_path}" | |
# Перевірка наявності маркера валідності | |
valid_marker = indices_path / "indices.valid" | |
if not valid_marker.exists(): | |
return False, f"Маркер валідності індексів не знайдено в {indices_path}" | |
# Перевірка наявності файлів індексів | |
required_files = ["docstore.json"] | |
for file in required_files: | |
if not (indices_path / file).exists(): | |
return False, f"Файл {file} не знайдено в {indices_path}" | |
# Перевірка наявності BM25 індексу | |
bm25_path = indices_path / "bm25" | |
if not bm25_path.exists() or not bm25_path.is_dir(): | |
return False, f"Директорія з BM25 індексом не знайдено в {indices_path}" | |
return True, "" | |
except Exception as e: | |
return False, f"Помилка при перевірці цілісності індексів: {str(e)}" | |
def check_indexing_availability(indices_path=None): | |
""" | |
Перевіряє доступність функціональності індексування. | |
Returns: | |
bool: True, якщо функціональність доступна, False - інакше | |
""" | |
try: | |
# Перевіряємо наявність необхідних модулів | |
import importlib | |
# Список необхідних модулів | |
required_modules = [ | |
"llama_index.core", | |
"llama_index.retrievers.bm25", | |
"llama_index.vector_stores.faiss", | |
"llama_index.embeddings.huggingface" | |
] | |
# Додаємо Google Embeddings до списку, якщо встановлено змінну середовища | |
if os.getenv("GEMINI_API_KEY"): | |
required_modules.append("google.genai") | |
# Перевіряємо кожен модуль | |
for module_name in required_modules: | |
try: | |
importlib.import_module(module_name) | |
except ImportError: | |
logger.warning(f"Модуль {module_name} не знайдено") | |
return False | |
# Всі модулі доступні | |
logger.info("Всі необхідні модулі для індексування доступні") | |
return True | |
except Exception as e: | |
logger.error(f"Помилка при перевірці доступності індексування: {e}") | |
return False | |
def validate_index_directory(indices_path): | |
""" | |
Перевіряє, чи директорія з індексами існує та містить необхідні файли. | |
Args: | |
indices_path (str): Шлях до директорії з індексами | |
Returns: | |
bool: True, якщо директорія валідна, False - інакше | |
""" | |
try: | |
from pathlib import Path | |
indices_path = Path(indices_path) | |
# Перевірка наявності директорії | |
if not indices_path.exists() or not indices_path.is_dir(): | |
return False | |
# Перевірка наявності необхідних файлів | |
required_files = ["docstore.json"] | |
for file in required_files: | |
if not (indices_path / file).exists(): | |
return False | |
return True | |
except Exception as e: | |
logger.error(f"Помилка при валідації директорії індексів: {str(e)}") | |
return False | |
def test_google_embeddings(): | |
""" | |
Функція для тестування та відлагодження Google Embeddings API. | |
Можна запустити як окремий скрипт для перевірки роботи API. | |
Запуск з командного рядка: | |
python -c "from modules.data_management.index_utils import test_google_embeddings; test_google_embeddings()" | |
""" | |
import os | |
import logging | |
# Налаштування логування | |
logging.basicConfig(level=logging.INFO) | |
logger = logging.getLogger(__name__) | |
logger.info("Тестування Google Embeddings API...") | |
# Отримання API ключа | |
api_key = os.getenv("GEMINI_API_KEY") | |
if not api_key: | |
logger.error("GEMINI_API_KEY не знайдений. Перевірте ваш .env файл або змінні середовища.") | |
return False | |
logger.info(f"API ключ Google знайдено: {api_key[:5]}...{api_key[-5:] if len(api_key) > 10 else '***'}") | |
try: | |
from google import genai | |
# Ініціалізація клієнта | |
genai.configure(api_key=api_key) | |
client = genai.Client() | |
logger.info("Google GenAI клієнт успішно ініціалізовано") | |
# Спроба отримати ембединги | |
text = ["Тестовий текст українською мовою"] | |
model = "text-embedding-004" | |
logger.info(f"Запит до моделі {model} з текстом: {text}") | |
result = client.models.embed_content( | |
model=model, | |
contents=text, | |
config={"task_type": "retrieval_query"} | |
) | |
# Отримання ембедингів | |
[embedding] = result.embeddings | |
embedding_values = embedding.values | |
logger.info(f"Ембединг успішно отримано, розмірність: {len(embedding_values)}") | |
logger.info(f"Перші 5 значень: {embedding_values[:5]}") | |
return True | |
except ImportError: | |
logger.error("Модуль google.genai не знайдено. Будь ласка, встановіть його: pip install google-genai") | |
return False | |
except Exception as e: | |
import traceback | |
logger.error(f"Помилка при тестуванні Google Embeddings API: {e}") | |
logger.error(traceback.format_exc()) | |
return False | |
if __name__ == "__main__": | |
test_google_embeddings() |