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()