543 lines
22 KiB
Python
543 lines
22 KiB
Python
"""
|
||
HTTP клиент для взаимодействия с внешним RAG сервисом.
|
||
|
||
Использует REST API для получения скоров и отправки примеров.
|
||
"""
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import httpx
|
||
|
||
from helper_bot.utils.metrics import track_errors, track_time
|
||
from logs.custom_logger import logger
|
||
|
||
from .base import ScoringResult
|
||
from .exceptions import InsufficientExamplesError, ScoringError, TextTooShortError
|
||
|
||
|
||
@dataclass
|
||
class SimilarPost:
|
||
"""Данные о похожем посте."""
|
||
|
||
similarity: float
|
||
created_at: int
|
||
post_id: Optional[int]
|
||
text: str
|
||
rag_score: Optional[float]
|
||
|
||
|
||
@dataclass
|
||
class SimilarPostsResult:
|
||
"""Результат поиска похожих постов."""
|
||
|
||
similar_count: int
|
||
similar_posts: List[SimilarPost]
|
||
max_similarity: float = 0.0
|
||
|
||
def __post_init__(self):
|
||
if self.similar_posts:
|
||
self.max_similarity = max(p.similarity for p in self.similar_posts)
|
||
|
||
|
||
class RagApiClient:
|
||
"""
|
||
HTTP клиент для взаимодействия с внешним RAG сервисом.
|
||
|
||
Использует REST API для:
|
||
- Получения скоров постов (POST /api/v1/score)
|
||
- Отправки положительных примеров (POST /api/v1/examples/positive)
|
||
- Отправки отрицательных примеров (POST /api/v1/examples/negative)
|
||
- Получения статистики (GET /api/v1/stats)
|
||
|
||
Attributes:
|
||
api_url: Базовый URL API сервиса
|
||
api_key: API ключ для аутентификации
|
||
timeout: Таймаут запросов в секундах
|
||
test_mode: Включен ли тестовый режим (добавляет заголовок X-Test-Mode: true)
|
||
enabled: Включен ли клиент
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_url: str,
|
||
api_key: str,
|
||
timeout: int = 30,
|
||
test_mode: bool = False,
|
||
enabled: bool = True,
|
||
):
|
||
"""
|
||
Инициализация клиента.
|
||
|
||
Args:
|
||
api_url: Базовый URL API (например, http://хх.ххх.ххх.хх/api/v1)
|
||
api_key: API ключ для аутентификации
|
||
timeout: Таймаут запросов в секундах
|
||
test_mode: Включен ли тестовый режим (добавляет заголовок X-Test-Mode: true к запросам examples)
|
||
enabled: Включен ли клиент
|
||
"""
|
||
# Убираем trailing slash если есть
|
||
self.api_url = api_url.rstrip("/")
|
||
self.api_key = api_key
|
||
self.timeout = timeout
|
||
self.test_mode = test_mode
|
||
self._enabled = enabled
|
||
|
||
# Создаем HTTP клиент
|
||
self._client = httpx.AsyncClient(
|
||
timeout=httpx.Timeout(timeout),
|
||
headers={
|
||
"X-API-Key": api_key,
|
||
"Content-Type": "application/json",
|
||
},
|
||
)
|
||
|
||
logger.info(
|
||
f"RagApiClient инициализирован (url={self.api_url}, enabled={enabled})"
|
||
)
|
||
|
||
@property
|
||
def source_name(self) -> str:
|
||
"""Имя источника для результатов."""
|
||
return "rag"
|
||
|
||
@property
|
||
def is_enabled(self) -> bool:
|
||
"""Проверяет, включен ли клиент."""
|
||
return self._enabled
|
||
|
||
async def close(self) -> None:
|
||
"""Закрывает HTTP клиент."""
|
||
await self._client.aclose()
|
||
|
||
@track_time("calculate_score", "rag_client")
|
||
@track_errors("rag_client", "calculate_score")
|
||
async def calculate_score(self, text: str) -> ScoringResult:
|
||
"""
|
||
Рассчитывает скор для текста поста через API.
|
||
|
||
Args:
|
||
text: Текст поста для оценки
|
||
|
||
Returns:
|
||
ScoringResult с оценкой
|
||
|
||
Raises:
|
||
ScoringError: При ошибке расчета
|
||
InsufficientExamplesError: Если недостаточно примеров
|
||
TextTooShortError: Если текст слишком короткий
|
||
"""
|
||
if not self._enabled:
|
||
raise ScoringError("RAG API клиент отключен")
|
||
|
||
if not text or not text.strip():
|
||
raise TextTooShortError("Текст пустой")
|
||
|
||
try:
|
||
response = await self._client.post(
|
||
f"{self.api_url}/score", json={"text": text.strip()}
|
||
)
|
||
|
||
# Обрабатываем различные статусы
|
||
if response.status_code == 400:
|
||
try:
|
||
error_data = response.json()
|
||
error_msg = error_data.get("detail", "Неизвестная ошибка")
|
||
except Exception:
|
||
error_msg = response.text or "Неизвестная ошибка"
|
||
|
||
logger.warning(f"RagApiClient: Ошибка валидации запроса: {error_msg}")
|
||
|
||
if (
|
||
"недостаточно" in error_msg.lower()
|
||
or "insufficient" in error_msg.lower()
|
||
):
|
||
raise InsufficientExamplesError(error_msg)
|
||
if "коротк" in error_msg.lower() or "short" in error_msg.lower():
|
||
raise TextTooShortError(error_msg)
|
||
raise ScoringError(f"Ошибка валидации: {error_msg}")
|
||
|
||
if response.status_code == 401:
|
||
logger.error("RagApiClient: Ошибка аутентификации: неверный API ключ")
|
||
raise ScoringError("Ошибка аутентификации: неверный API ключ")
|
||
|
||
if response.status_code == 404:
|
||
logger.error("RagApiClient: RAG API endpoint не найден")
|
||
raise ScoringError("RAG API endpoint не найден")
|
||
|
||
if response.status_code >= 500:
|
||
logger.error(
|
||
f"RagApiClient: Ошибка сервера RAG API: {response.status_code}"
|
||
)
|
||
raise ScoringError(f"Ошибка сервера RAG API: {response.status_code}")
|
||
|
||
# Проверяем успешный статус
|
||
if response.status_code != 200:
|
||
response.raise_for_status()
|
||
|
||
data = response.json()
|
||
|
||
# Парсим ответ
|
||
score = float(data.get("rag_score", 0.0))
|
||
confidence = (
|
||
float(data.get("rag_confidence", 0.0))
|
||
if data.get("rag_confidence") is not None
|
||
else None
|
||
)
|
||
rag_score_pos_only_raw = data.get("rag_score_pos_only")
|
||
rag_score_pos_only = (
|
||
float(rag_score_pos_only_raw)
|
||
if rag_score_pos_only_raw is not None
|
||
else None
|
||
)
|
||
|
||
# Форматируем confidence для логирования
|
||
confidence_str = f"{confidence:.4f}" if confidence is not None else "None"
|
||
rag_score_pos_only_str = (
|
||
f"{rag_score_pos_only:.4f}"
|
||
if rag_score_pos_only is not None
|
||
else "None"
|
||
)
|
||
|
||
logger.info(
|
||
f"RagApiClient: Скор успешно получен из API - "
|
||
f"rag_score={score:.4f} (type: {type(score).__name__}), "
|
||
f"rag_confidence={confidence_str}, "
|
||
f"rag_score_pos_only={rag_score_pos_only_str}, "
|
||
f"raw_response_rag_score={data.get('rag_score')}, "
|
||
f"raw_response_rag_score_pos_only={rag_score_pos_only_raw}"
|
||
)
|
||
|
||
return ScoringResult(
|
||
score=score,
|
||
source=self.source_name,
|
||
model=data.get("meta", {}).get("model", "rag-service"),
|
||
confidence=confidence,
|
||
metadata={
|
||
"rag_score_pos_only": (
|
||
float(data.get("rag_score_pos_only", 0.0))
|
||
if data.get("rag_score_pos_only") is not None
|
||
else None
|
||
),
|
||
"positive_examples": data.get("meta", {}).get("positive_examples"),
|
||
"negative_examples": data.get("meta", {}).get("negative_examples"),
|
||
},
|
||
)
|
||
|
||
except httpx.TimeoutException:
|
||
logger.error(f"RagApiClient: Таймаут запроса к RAG API (>{self.timeout}с)")
|
||
raise ScoringError(f"Таймаут запроса к RAG API (>{self.timeout}с)")
|
||
except httpx.RequestError as e:
|
||
logger.error(f"RagApiClient: Ошибка подключения к RAG API: {e}")
|
||
raise ScoringError(f"Ошибка подключения к RAG API: {e}")
|
||
except (KeyError, ValueError, TypeError) as e:
|
||
logger.error(
|
||
f"RagApiClient: Ошибка парсинга ответа: {e}, response: {response.text if 'response' in locals() else 'N/A'}"
|
||
)
|
||
raise ScoringError(f"Ошибка парсинга ответа от RAG API: {e}")
|
||
except InsufficientExamplesError:
|
||
raise
|
||
except TextTooShortError:
|
||
raise
|
||
except ScoringError:
|
||
# Уже залогированные ошибки (401, 404, 500, таймауты и т.д.) - просто пробрасываем
|
||
raise
|
||
except Exception as e:
|
||
# Только действительно неожиданные ошибки логируем здесь
|
||
logger.error(
|
||
f"RagApiClient: Неожиданная ошибка при расчете скора: {e}",
|
||
exc_info=True,
|
||
)
|
||
raise ScoringError(f"Неожиданная ошибка: {e}")
|
||
|
||
@track_time("add_positive_example", "rag_client")
|
||
async def add_positive_example(self, text: str) -> None:
|
||
"""
|
||
Добавляет текст как положительный пример (опубликованный пост).
|
||
|
||
Args:
|
||
text: Текст опубликованного поста
|
||
"""
|
||
if not self._enabled:
|
||
return
|
||
|
||
if not text or not text.strip():
|
||
return
|
||
|
||
try:
|
||
# Формируем заголовки (добавляем X-Test-Mode если включен тестовый режим)
|
||
headers = {}
|
||
if self.test_mode:
|
||
headers["X-Test-Mode"] = "true"
|
||
|
||
response = await self._client.post(
|
||
f"{self.api_url}/examples/positive",
|
||
json={"text": text.strip()},
|
||
headers=headers,
|
||
)
|
||
|
||
if response.status_code == 200 or response.status_code == 201:
|
||
logger.info("RagApiClient: Положительный пример успешно добавлен")
|
||
elif response.status_code == 400:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка валидации при добавлении положительного примера: {response.text}"
|
||
)
|
||
else:
|
||
logger.warning(
|
||
f"RagApiClient: Неожиданный статус при добавлении положительного примера: {response.status_code}"
|
||
)
|
||
|
||
except httpx.TimeoutException:
|
||
logger.warning(
|
||
f"RagApiClient: Таймаут при добавлении положительного примера"
|
||
)
|
||
except httpx.RequestError as e:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка подключения при добавлении положительного примера: {e}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"RagApiClient: Ошибка добавления положительного примера: {e}")
|
||
|
||
@track_time("add_negative_example", "rag_client")
|
||
async def add_negative_example(self, text: str) -> None:
|
||
"""
|
||
Добавляет текст как отрицательный пример (отклоненный пост).
|
||
|
||
Args:
|
||
text: Текст отклоненного поста
|
||
"""
|
||
if not self._enabled:
|
||
return
|
||
|
||
if not text or not text.strip():
|
||
return
|
||
|
||
try:
|
||
# Формируем заголовки (добавляем X-Test-Mode если включен тестовый режим)
|
||
headers = {}
|
||
if self.test_mode:
|
||
headers["X-Test-Mode"] = "true"
|
||
|
||
response = await self._client.post(
|
||
f"{self.api_url}/examples/negative",
|
||
json={"text": text.strip()},
|
||
headers=headers,
|
||
)
|
||
|
||
if response.status_code == 200 or response.status_code == 201:
|
||
logger.info("RagApiClient: Отрицательный пример успешно добавлен")
|
||
elif response.status_code == 400:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка валидации при добавлении отрицательного примера: {response.text}"
|
||
)
|
||
else:
|
||
logger.warning(
|
||
f"RagApiClient: Неожиданный статус при добавлении отрицательного примера: {response.status_code}"
|
||
)
|
||
|
||
except httpx.TimeoutException:
|
||
logger.warning(
|
||
f"RagApiClient: Таймаут при добавлении отрицательного примера"
|
||
)
|
||
except httpx.RequestError as e:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка подключения при добавлении отрицательного примера: {e}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"RagApiClient: Ошибка добавления отрицательного примера: {e}")
|
||
|
||
async def get_stats(self) -> Dict[str, Any]:
|
||
"""
|
||
Получает статистику от RAG API через endpoint /stats.
|
||
|
||
Returns:
|
||
Словарь со статистикой или пустой словарь при ошибке
|
||
"""
|
||
if not self._enabled:
|
||
logger.debug("RagApiClient: get_stats пропущен - клиент отключен")
|
||
return {}
|
||
|
||
try:
|
||
logger.debug(f"RagApiClient: Запрос статистики от {self.api_url}/stats")
|
||
response = await self._client.get(f"{self.api_url}/stats")
|
||
|
||
if response.status_code == 200:
|
||
data = response.json()
|
||
logger.info(
|
||
f"RagApiClient: Статистика получена успешно: "
|
||
f"model_loaded={data.get('model_loaded')}, "
|
||
f"model_name={data.get('model_name')}, "
|
||
f"vector_store={data.get('vector_store', {}).get('total_count', 'N/A')} примеров"
|
||
)
|
||
return data
|
||
elif response.status_code == 401 or response.status_code == 403:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка авторизации при получении статистики: "
|
||
f"status={response.status_code}, body={response.text[:200]}"
|
||
)
|
||
return {}
|
||
else:
|
||
logger.warning(
|
||
f"RagApiClient: Неожиданный статус при получении статистики: "
|
||
f"status={response.status_code}, body={response.text[:200]}"
|
||
)
|
||
return {}
|
||
|
||
except httpx.TimeoutException:
|
||
logger.warning(
|
||
f"RagApiClient: Таймаут при получении статистики (timeout={self.timeout}s)"
|
||
)
|
||
return {}
|
||
except httpx.RequestError as e:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка подключения при получении статистики: {e}"
|
||
)
|
||
return {}
|
||
except Exception as e:
|
||
logger.error(f"RagApiClient: Ошибка получения статистики: {e}")
|
||
return {}
|
||
|
||
def get_stats_sync(self) -> Dict[str, Any]:
|
||
"""
|
||
Синхронная версия get_stats для использования в get_stats() ScoringManager.
|
||
|
||
Внимание: Это заглушка, реальная статистика будет получена асинхронно.
|
||
"""
|
||
return {
|
||
"enabled": self._enabled,
|
||
"api_url": self.api_url,
|
||
"timeout": self.timeout,
|
||
}
|
||
|
||
@track_time("find_similar_posts", "rag_client")
|
||
async def find_similar_posts(
|
||
self, text: str, threshold: float = 0.9, hours: int = 24
|
||
) -> Optional[SimilarPostsResult]:
|
||
"""
|
||
Ищет похожие посты за последние N часов.
|
||
|
||
Args:
|
||
text: Текст поста для поиска похожих
|
||
threshold: Порог схожести (0.0-1.0), по умолчанию 0.9
|
||
hours: За сколько часов искать (1-168), по умолчанию 24
|
||
|
||
Returns:
|
||
SimilarPostsResult с информацией о похожих постах или None при ошибке
|
||
"""
|
||
if not self._enabled:
|
||
return None
|
||
|
||
if not text or not text.strip():
|
||
return None
|
||
|
||
try:
|
||
response = await self._client.post(
|
||
f"{self.api_url}/similar",
|
||
json={"text": text.strip(), "threshold": threshold, "hours": hours},
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
data = response.json()
|
||
similar_posts = []
|
||
|
||
for post_data in data.get("similar_posts", []):
|
||
similar_posts.append(
|
||
SimilarPost(
|
||
similarity=float(post_data.get("similarity", 0.0)),
|
||
created_at=int(post_data.get("created_at", 0)),
|
||
post_id=post_data.get("post_id"),
|
||
text=post_data.get("text", ""),
|
||
rag_score=post_data.get("rag_score"),
|
||
)
|
||
)
|
||
|
||
result = SimilarPostsResult(
|
||
similar_count=data.get("similar_count", 0),
|
||
similar_posts=similar_posts,
|
||
)
|
||
|
||
if result.similar_count > 0:
|
||
logger.info(
|
||
f"RagApiClient: Найдено {result.similar_count} похожих постов "
|
||
f"(max_similarity={result.max_similarity:.2%})"
|
||
)
|
||
|
||
return result
|
||
else:
|
||
logger.warning(
|
||
f"RagApiClient: Неожиданный статус при поиске похожих постов: "
|
||
f"{response.status_code}, body: {response.text}"
|
||
)
|
||
return None
|
||
|
||
except httpx.TimeoutException:
|
||
logger.warning("RagApiClient: Таймаут при поиске похожих постов")
|
||
return None
|
||
except httpx.RequestError as e:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка подключения при поиске похожих постов: {e}"
|
||
)
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"RagApiClient: Ошибка поиска похожих постов: {e}")
|
||
return None
|
||
|
||
@track_time("add_submitted_post", "rag_client")
|
||
async def add_submitted_post(
|
||
self, text: str, post_id: Optional[int] = None, rag_score: Optional[float] = None
|
||
) -> bool:
|
||
"""
|
||
Добавляет пост в коллекцию submitted для поиска похожих.
|
||
|
||
Args:
|
||
text: Текст поста
|
||
post_id: ID поста (опционально)
|
||
rag_score: RAG скор на момент добавления (опционально)
|
||
|
||
Returns:
|
||
True если пост успешно добавлен
|
||
"""
|
||
if not self._enabled:
|
||
return False
|
||
|
||
if not text or not text.strip():
|
||
return False
|
||
|
||
try:
|
||
payload = {"text": text.strip()}
|
||
if post_id is not None:
|
||
payload["post_id"] = post_id
|
||
if rag_score is not None:
|
||
payload["rag_score"] = rag_score
|
||
|
||
response = await self._client.post(
|
||
f"{self.api_url}/submitted",
|
||
json=payload,
|
||
)
|
||
|
||
if response.status_code in (200, 201):
|
||
data = response.json()
|
||
logger.debug(
|
||
f"RagApiClient: Пост добавлен в submitted "
|
||
f"(post_id={post_id}, submitted_count={data.get('submitted_count', 'N/A')})"
|
||
)
|
||
return True
|
||
else:
|
||
logger.warning(
|
||
f"RagApiClient: Неожиданный статус при добавлении в submitted: "
|
||
f"{response.status_code}"
|
||
)
|
||
return False
|
||
|
||
except httpx.TimeoutException:
|
||
logger.warning("RagApiClient: Таймаут при добавлении в submitted")
|
||
return False
|
||
except httpx.RequestError as e:
|
||
logger.warning(
|
||
f"RagApiClient: Ошибка подключения при добавлении в submitted: {e}"
|
||
)
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"RagApiClient: Ошибка добавления в submitted: {e}")
|
||
return False
|