DocUA's picture
Єдиний коміт - очищення історії
4ad5efa
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()