umb/bot/utils/memory.py

import json
import logging
import math
import re

import aiohttp

from bot.utils.proxy import get_proxy_connector
from bot.utils.database import get_summaries, save_conversation_summary
from bot.utils.ai_client import get_client_timeout, ask_ai_simple
from config import ROUTERAI_API_KEY, ROUTERAI_BASE_URL

logger = logging.getLogger(__name__)

EMBEDDING_MODEL = "openai/text-embedding-3-small"
EMBEDDING_URL = f"{ROUTERAI_BASE_URL}/embeddings"


async def create_embedding(text: str) -> list[float] | None:
    if not ROUTERAI_API_KEY:
        logger.warning("Embedding skipped | ROUTERAI_API_KEY not set")
        return None

    headers = {
        "Authorization": f"Bearer {ROUTERAI_API_KEY}",
        "Content-Type": "application/json",
    }

    payload = {
        "model": EMBEDDING_MODEL,
        "input": text[:8000],
        "encoding_format": "float",
    }

    connector = get_proxy_connector()

    try:
        async with aiohttp.ClientSession(
            timeout=get_client_timeout(30), connector=connector
        ) as session:
            async with session.post(EMBEDDING_URL, json=payload, headers=headers) as response:
                if response.status != 200:
                    error_body = await response.text()
                    logger.warning(
                        "Embedding error | status=%s error=%s",
                        response.status,
                        error_body[:200],
                    )
                    return None

                data = await response.json()
                embedding = data["data"][0]["embedding"]
                logger.info(
                    "Embedding created | dim=%d input_len=%d",
                    len(embedding),
                    min(len(text), 8000),
                )
                return embedding
    except Exception as exc:
        logger.error("Embedding request error: %s", exc)
        return None
    finally:
        if connector:
            await connector.close()


def cosine_similarity(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(y * y for y in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


async def find_relevant_summaries(
    user_id: int, chat_id: int, query: str, top_k: int = 3
) -> list[str]:
    query_emb = await create_embedding(query)
    if not query_emb:
        return []

    summaries = await get_summaries(user_id, chat_id, limit=20)
    if not summaries:
        logger.info(
            "Relevant summaries | user=%d chat=%d no summaries found", user_id, chat_id
        )
        return []

    scored = []
    for s in summaries:
        if not s["embedding"]:
            continue
        try:
            emb = json.loads(s["embedding"])
            score = cosine_similarity(query_emb, emb)
            scored.append((score, s["summary"]))
        except Exception:
            continue

    scored.sort(key=lambda x: x[0], reverse=True)
    top = [text for _, text in scored[:top_k]]
    top_score = scored[0][0] if scored else 0

    logger.info(
        "Relevant summaries | user=%d chat=%d found=%d top_score=%.3f",
        user_id,
        chat_id,
        len(top),
        top_score,
    )
    return top


async def save_summary_with_embedding(
    user_id: int, chat_id: int, summary_text: str
) -> None:
    emb = await create_embedding(summary_text)
    embedding_json = json.dumps(emb) if emb else None
    await save_conversation_summary(user_id, chat_id, summary_text, embedding_json)
    logger.info(
        "Summary saved | user=%d chat=%d summary_len=%d emb=%s",
        user_id,
        chat_id,
        len(summary_text),
        "yes" if emb else "no",
    )


async def generate_summary(messages: list[dict]) -> str | None:
    messages_text = "\n".join(
        f"{'Пользователь' if m['role'] == 'user' else 'Астра'}: {m['content'][:300]}"
        for m in messages[-50:]
    )

    prompt = (
        "Сделай краткую выжимку этого диалога (3-5 предложений). "
        "Выдели ключевые темы, факты и предпочтения пользователя:\n\n"
        f"{messages_text}"
    )

    try:
        summary = await ask_ai_simple(prompt)
        if summary and len(summary) > 20:
            cleaned = re.sub(r"<[^>]+>", "", summary)
            cleaned = cleaned.strip()
            logger.info("Summary generated | len=%d", len(cleaned))
            return cleaned
        elif summary:
            logger.warning("Summary too short | len=%d", len(summary))
        else:
            logger.warning("Summary generation returned None")
    except Exception as exc:
        logger.error("Summary generation error: %s", exc)

    return None