kerrad/rag-service
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add submitted collection, /similar and /submitted endpoints (Stage 4)

Browse Source 
Made-with: Cursor
This commit is contained in:
Andrey
2026-02-28 19:00:22 +03:00
parent 955f518429
commit a1d6d2d860
15 changed files with 1308 additions and 400 deletions
Download Patch File Download Diff File Expand all files Collapse all files

283
app/services/rag_service.py
View File

@@ -9,7 +9,7 @@ import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional
from typing import Any
import numpy as np
@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
class ScoringResult:
"""
Результат оценки поста.
Attributes:
score: Оценка от 0.0 до 1.0 (вероятность публикации)
confidence: Уверенность в оценке
@@ -39,6 +39,7 @@ class ScoringResult:
model: Название используемой модели
timestamp: Время получения оценки
"""
score: float
confidence: float
score_pos_only: float
@@ -46,8 +47,8 @@ class ScoringResult:
negative_examples: int
model: str
timestamp: int = field(default_factory=lambda: int(datetime.now().timestamp()))
def to_dict(self) -> Dict[str, Any]:
def to_dict(self) -> dict[str, Any]:
"""Преобразует результат в словарь."""
return {
"rag_score": round(self.score, 4),
@@ -58,31 +59,31 @@ class ScoringResult:
"negative_examples": self.negative_examples,
"model": self.model,
"timestamp": self.timestamp,
}
},
}
class RAGService:
"""
RAG сервис для оценки постов на основе векторного сходства.
Использует sentence-transformers для создания эмбеддингов текста и сравнивает
их с эталонными примерами (опубликованные vs отклоненные посты).
Attributes:
model_name: Название модели HuggingFace
vector_store: Хранилище векторов
min_text_length: Минимальная длина текста для обработки
"""
def __init__(
self,
settings: Optional[Settings] = None,
vector_store: Optional[VectorStore] = None,
settings: Settings | None = None,
vector_store: VectorStore | None = None,
):
"""
Инициализация RAG сервиса.
Args:
settings: Настройки сервиса (берутся из get_settings() если не переданы)
vector_store: Хранилище векторов (создается автоматически если не передано)
@@ -91,96 +92,102 @@ class RAGService:
self.model_name = self._settings.model_name
self.cache_dir = self._settings.cache_dir
self.min_text_length = self._settings.min_text_length
# Модель загружается лениво
self._model = None
self._device = None
self._model_loaded = False
# Хранилище векторов
self.vector_store = vector_store or VectorStore(
vector_dim=self._settings.vector_dim,
max_examples=self._settings.max_examples,
max_submitted=self._settings.max_submitted,
storage_path=self._settings.vectors_path,
submitted_path=self._settings.submitted_path,
score_multiplier=self._settings.score_multiplier,
k=3, # Фиксированное значение k для топ-k ближайших примеров
)
logger.info(f"RAGService инициализирован (model={self.model_name})")
@property
def is_model_loaded(self) -> bool:
"""Проверяет, загружена ли модель."""
return self._model_loaded
async def load_model(self) -> None:
"""
Загружает модель и токенизатор.
Выполняется асинхронно в отдельном потоке чтобы не блокировать event loop.
"""
if self._model_loaded:
return
logger.info(f"RAGService: Загрузка модели {self.model_name}...")
try:
# Загрузка в отдельном потоке
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, self._load_model_sync)
self._model_loaded = True
logger.info(f"RAGService: Модель {self.model_name} успешно загружена")
except Exception as e:
logger.error(f"RAGService: Ошибка загрузки модели: {e}")
raise ModelNotLoadedError(f"Не удалось загрузить модель {self.model_name}: {e}")
def _load_model_sync(self) -> None:
"""Синхронная загрузка модели (вызывается в executor)."""
logger.info("RAGService: Начало _load_model_sync, импорт sentence_transformers...")
from sentence_transformers import SentenceTransformer
import torch
from sentence_transformers import SentenceTransformer
# Определяем устройство
self._device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"RAGService: Устройство определено: {self._device}")
# Загружаем модель SentenceTransformer
logger.info(f"RAGService: Загрузка модели из {self.model_name} (это может занять несколько минут)...")
logger.info(
f"RAGService: Загрузка модели из {self.model_name} (это может занять несколько минут)..."
)
self._model = SentenceTransformer(
self.model_name,
cache_folder=self.cache_dir,
device=self._device,
)
logger.info(f"RAGService: Модель готова на устройстве: {self._device}")
def _get_embedding_sync(self, text: str) -> np.ndarray:
"""
Получает эмбеддинг текста (синхронно).
Использует SentenceTransformer для получения нормализованного эмбеддинга.
Args:
text: Текст для векторизации
Returns:
Numpy массив с эмбеддингом (384 измерений для all-MiniLM-L12-v2)
"""
# SentenceTransformer автоматически нормализует эмбеддинги
embedding = self._model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
return embedding.flatten()
def _get_embeddings_batch_sync(self, texts: List[str], batch_size: int = 16) -> List[np.ndarray]:
def _get_embeddings_batch_sync(
self, texts: list[str], batch_size: int = 16
) -> list[np.ndarray]:
"""
Получает эмбеддинги для батча текстов (синхронно).
Обрабатывает тексты пачками для эффективного использования GPU/CPU.
Args:
texts: Список текстов для векторизации
batch_size: Размер батча
Returns:
Список numpy массивов с эмбеддингами
"""
@@ -192,32 +199,34 @@ class RAGService:
normalize_embeddings=True,
show_progress_bar=False,
)
# Преобразуем в список отдельных массивов
return [emb.flatten() for emb in embeddings]
async def get_embeddings_batch(self, texts: List[str], batch_size: Optional[int] = None) -> List[np.ndarray]:
async def get_embeddings_batch(
self, texts: list[str], batch_size: int | None = None
) -> list[np.ndarray]:
"""
Получает эмбеддинги для батча текстов (асинхронно).
Args:
texts: Список текстов для векторизации
batch_size: Размер батча (берется из настроек если не указан)
Returns:
Список numpy массивов с эмбеддингами
"""
if not self._model_loaded:
await self.load_model()
if not self._model_loaded:
raise ModelNotLoadedError("Модель не загружена")
batch_size = batch_size or self._settings.batch_size
# Очищаем тексты
clean_texts = [self._clean_text(text) for text in texts]
# Выполняем батч-обработку в thread pool
loop = asyncio.get_event_loop()
embeddings = await loop.run_in_executor(
@@ -226,71 +235,67 @@ class RAGService:
clean_texts,
batch_size,
)
return embeddings
async def get_embedding(self, text: str) -> np.ndarray:
"""
Получает эмбеддинг текста (асинхронно).
Args:
text: Текст для векторизации
Returns:
Numpy массив с эмбеддингом
Raises:
ModelNotLoadedError: Если модель не загружена
TextTooShortError: Если текст слишком короткий
"""
if not self._model_loaded:
await self.load_model()
if not self._model_loaded:
raise ModelNotLoadedError("Модель не загружена")
# Очищаем текст
clean_text = self._clean_text(text)
if len(clean_text) < self.min_text_length:
raise TextTooShortError(
f"Текст слишком короткий (минимум {self.min_text_length} символов)"
)
# Выполняем в отдельном потоке
loop = asyncio.get_event_loop()
embedding = await loop.run_in_executor(
None,
self._get_embedding_sync,
clean_text
)
embedding = await loop.run_in_executor(None, self._get_embedding_sync, clean_text)
return embedding
def _clean_text(self, text: str) -> str:
"""Очищает текст от лишних символов."""
if not text:
return ""
# Удаляем лишние пробелы и переносы строк
clean = " ".join(text.split())
# Удаляем служебные символы (например "^" для helper сообщений)
if clean == "^":
return ""
return clean.strip()
async def calculate_score(self, text: str) -> ScoringResult:
"""
Рассчитывает скор для текста поста.
Args:
text: Текст поста для оценки
Returns:
ScoringResult с оценкой
Raises:
ScoringError: При ошибке расчета
InsufficientExamplesError: Если недостаточно примеров
@@ -299,16 +304,17 @@ class RAGService:
try:
# Получаем эмбеддинг текста
embedding = await self.get_embedding(text)
# Логируем первые элементы вектора для отладки
logger.debug(
f"RAGService: embedding[:3]={embedding[:3].tolist()}, "
f"text_preview='{text[:30]}'"
f"RAGService: embedding[:3]={embedding[:3].tolist()}, text_preview='{text[:30]}'"
)
# Рассчитываем скор через VectorStore
score, confidence, score_pos_only = self.vector_store.calculate_similarity_score(embedding)
score, confidence, score_pos_only = self.vector_store.calculate_similarity_score(
embedding
)
return ScoringResult(
score=score,
confidence=confidence,
@@ -317,22 +323,22 @@ class RAGService:
negative_examples=self.vector_store.negative_count,
model=self.model_name,
)
except (InsufficientExamplesError, TextTooShortError):
# Пробрасываем ожидаемые исключения
raise
except Exception as e:
logger.error(f"RAGService: Ошибка расчета скора: {e}")
raise ScoringError(f"Ошибка расчета скора: {e}")
async def add_positive_example(self, text: str) -> bool:
"""
Добавляет текст как положительный пример (опубликованный пост).
Args:
text: Текст опубликованного поста
Returns:
True если пример добавлен, False если дубликат/короткий текст
"""
@@ -341,32 +347,32 @@ class RAGService:
if len(clean_text) < self.min_text_length:
logger.debug("RAGService: Текст слишком короткий для примера, пропускаем")
return False
# Получаем эмбеддинг
embedding = await self.get_embedding(clean_text)
# Вычисляем хеш для дедупликации
text_hash = VectorStore.compute_text_hash(clean_text)
# Добавляем в хранилище
added = self.vector_store.add_positive(embedding, text_hash)
if added:
logger.info("RAGService: Добавлен положительный пример")
return added
except Exception as e:
logger.error(f"RAGService: Ошибка добавления положительного примера: {e}")
return False
async def add_negative_example(self, text: str) -> bool:
"""
Добавляет текст как отрицательный пример (отклоненный пост).
Args:
text: Текст отклоненного поста
Returns:
True если пример добавлен, False если дубликат/короткий текст
"""
@@ -375,29 +381,102 @@ class RAGService:
if len(clean_text) < self.min_text_length:
logger.debug("RAGService: Текст слишком короткий для примера, пропускаем")
return False
# Получаем эмбеддинг
embedding = await self.get_embedding(clean_text)
# Вычисляем хеш для дедупликации
text_hash = VectorStore.compute_text_hash(clean_text)
# Добавляем в хранилище
added = self.vector_store.add_negative(embedding, text_hash)
if added:
logger.info("RAGService: Добавлен отрицательный пример")
return added
except Exception as e:
logger.error(f"RAGService: Ошибка добавления отрицательного примера: {e}")
return False
async def add_submitted_post(
self,
text: str,
post_id: int | None = None,
rag_score: float | None = None,
) -> bool:
"""
Добавляет submitted-пост в коллекцию для индексации.
Args:
text: Текст поста
post_id: ID поста (опционально)
rag_score: RAG скор поста (опционально)
Returns:
True если добавлен, False если дубликат/короткий текст
"""
try:
clean_text = self._clean_text(text)
if len(clean_text) < self.min_text_length:
logger.debug("RAGService: Текст слишком короткий для submitted, пропускаем")
return False
embedding = await self.get_embedding(clean_text)
text_hash = VectorStore.compute_text_hash(clean_text)
created_at = int(datetime.now().timestamp())
added = self.vector_store.add_submitted(
vector=embedding,
text_hash=text_hash,
created_at=created_at,
post_id=post_id,
text=clean_text,
rag_score=rag_score,
)
if added:
logger.info("RAGService: Добавлен submitted-пост")
return added
except Exception as e:
logger.error(f"RAGService: Ошибка добавления submitted-поста: {e}")
return False
async def find_similar_posts(
self,
text: str,
threshold: float = 0.9,
hours: int = 24,
) -> list[dict[str, Any]]:
"""
Ищет похожие submitted-посты за последние N часов.
Args:
text: Текст для поиска
threshold: Минимальный порог similarity (0.0 - 1.0)
hours: Количество часов для фильтрации
Returns:
Список dict с полями: similarity, created_at, post_id, text, rag_score
"""
try:
embedding = await self.get_embedding(text)
return self.vector_store.find_similar_submitted(
vector=embedding,
threshold=threshold,
hours=hours,
)
except Exception as e:
logger.error(f"RAGService: Ошибка поиска похожих постов: {e}")
return []
async def warmup(self) -> bool:
"""
Прогревает модель (загружает если не загружена).
Returns:
True если модель загружена успешно
"""
@@ -407,13 +486,15 @@ class RAGService:
except Exception as e:
logger.error(f"RAGService: Ошибка прогрева модели: {e}")
return False
def save_vectors(self) -> None:
"""Сохраняет векторы на диск."""
"""Сохраняет векторы на диск (включая submitted)."""
if self.vector_store.storage_path:
self.vector_store.save_to_disk()
def get_stats(self) -> Dict[str, Any]:
if self.vector_store.submitted_path:
self.vector_store.save_submitted_to_disk()
def get_stats(self) -> dict[str, Any]:
"""Возвращает статистику сервиса."""
return {
"model_name": self.model_name,
@@ -424,13 +505,13 @@ class RAGService:
# Глобальный экземпляр сервиса (singleton)
_rag_service: Optional[RAGService] = None
_rag_service: RAGService | None = None
def get_rag_service() -> RAGService:
"""
Возвращает глобальный экземпляр RAG сервиса.
Returns:
RAGService: Экземпляр сервиса
"""