bot/utils/embedding_utils.py - master - Em1Te/digital-PMPO

bot/utils/embedding_utils.py
194 строки7 KB
Em1Te
test
14 май 2026, 19:40
14 май 2026, 19:40fd61bcf
import asyncio
import logging
import time
from typing import Optional

import aiohttp

from bot.config import (
    RAG_DEFAULT_EMBEDDING_MODEL,
    RAG_EMBEDDING_DIMENSION,
    RAG_EMBEDDING_MAX_CHARS,
    RAG_EMBEDDING_TIMEOUT_SEC,
)
from bot.utils.openrouter_key_manager import OpenRouterKeyManager


class EmbeddingProcessor:
    """Embeddings client over OpenRouter with key rotation and retries."""

    def __init__(self):
        self._session: aiohttp.ClientSession | None = None
        self._timeout = aiohttp.ClientTimeout(total=RAG_EMBEDDING_TIMEOUT_SEC)
        self._key_manager = OpenRouterKeyManager()
        self._url = "https://openrouter.ai/api/v1/embeddings"
        self._model = RAG_DEFAULT_EMBEDDING_MODEL
        self._dimensions = RAG_EMBEDDING_DIMENSION
        self._max_chars = max(1000, RAG_EMBEDDING_MAX_CHARS)

        if self._key_manager.has_available_keys():
            logging.info("Embedding processor is ready")
        else:
            logging.warning("No OpenRouter keys found for embeddings")

    def is_available(self) -> bool:
        return self._key_manager.has_available_keys()

    async def _ensure_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(timeout=self._timeout, trust_env=True)
        return self._session

    async def _write_usage(
        self,
        *,
        event: str,
        team_id: int | None,
        employee_tg_id: int | None,
        model: str,
        input_tokens: int,
        total_tokens: int,
        duration_ms: int,
        attempts: int,
        response_text: str | None = None,
    ) -> None:
        try:
            from bot.core.database import db_add_llm_token_usage

            await db_add_llm_token_usage(
                event=event,
                team_id=team_id,
                employee_tg_id=employee_tg_id,
                provider="openrouter",
                model=model,
                model_version=None,
                input_tokens=input_tokens,
                output_tokens=0,
                total_tokens=total_tokens,
                duration_ms=duration_ms,
                attempts=attempts,
                response_text=response_text,
            )
        except Exception as e:
            logging.error(f"Embedding usage logging failed: {e}")

    async def embed_text(
        self,
        text: str,
        *,
        team_id: int | None = None,
        employee_tg_id: int | None = None,
        usage_event: str = "rag_embedding",
    ) -> Optional[list[float]]:
        if not self.is_available():
            return None

        cleaned = (text or "").strip()
        if not cleaned:
            return None
        if len(cleaned) > self._max_chars:
            cleaned = cleaned[: self._max_chars]

        session = await self._ensure_session()
        start_ts = time.perf_counter()
        attempts = 0

        # Phase 1: free keys
        while True:
            key_model_pair = self._key_manager.get_free_key_and_model()
            if not key_model_pair:
                break

            api_key, _ = key_model_pair
            attempts += 1
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            }
            payload = {
                "model": self._model,
                "input": cleaned,
                "dimensions": self._dimensions,
            }

            try:
                async with session.post(self._url, json=payload, headers=headers) as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        embedding_data = (data.get("data") or [{}])[0]
                        embedding = embedding_data.get("embedding")
                        if isinstance(embedding, list) and embedding:
                            self._key_manager.mark_result(api_key, success=True)
                            duration_ms = int((time.perf_counter() - start_ts) * 1000)
                            usage = data.get("usage", {}) or {}
                            prompt_tokens = int(usage.get("prompt_tokens") or 0)
                            total_tokens = int(usage.get("total_tokens") or prompt_tokens)
                            await self._write_usage(
                                event=usage_event,
                                team_id=team_id,
                                employee_tg_id=employee_tg_id,
                                model=self._model,
                                input_tokens=prompt_tokens,
                                total_tokens=total_tokens,
                                duration_ms=duration_ms,
                                attempts=attempts,
                            )
                            return [float(v) for v in embedding]

                        self._key_manager.mark_result(api_key, success=False)
                        continue

                    if resp.status == 403:
                        self._key_manager.mark_result(api_key, success=False, is_403=True)
                    else:
                        self._key_manager.mark_result(api_key, success=False)
            except (aiohttp.ClientError, asyncio.TimeoutError):
                self._key_manager.mark_result(api_key, success=False)

        # Phase 2: paid key fallback (up to 3 retries)
        paid_key_pair = self._key_manager.get_paid_key_and_model()
        if not paid_key_pair:
            return None

        api_key, _ = paid_key_pair
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": self._model,
            "input": cleaned,
            "dimensions": self._dimensions,
        }

        for i in range(3):
            attempts += 1
            try:
                async with session.post(self._url, json=payload, headers=headers) as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        embedding_data = (data.get("data") or [{}])[0]
                        embedding = embedding_data.get("embedding")
                        if isinstance(embedding, list) and embedding:
                            duration_ms = int((time.perf_counter() - start_ts) * 1000)
                            usage = data.get("usage", {}) or {}
                            prompt_tokens = int(usage.get("prompt_tokens") or 0)
                            total_tokens = int(usage.get("total_tokens") or prompt_tokens)
                            await self._write_usage(
                                event=usage_event,
                                team_id=team_id,
                                employee_tg_id=employee_tg_id,
                                model=self._model,
                                input_tokens=prompt_tokens,
                                total_tokens=total_tokens,
                                duration_ms=duration_ms,
                                attempts=attempts,
                            )
                            return [float(v) for v in embedding]
            except (aiohttp.ClientError, asyncio.TimeoutError):
                pass
            await asyncio.sleep(1 + i)
        return None


embedding_processor = EmbeddingProcessor()