134 lines
4.4 KiB
Python
134 lines
4.4 KiB
Python
import json
|
|
import logging
|
|
import math
|
|
import re
|
|
|
|
from bot.utils.ai_client import _get_connector
|
|
from bot.utils.database import get_summaries, save_conversation_summary
|
|
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",
|
|
}
|
|
|
|
import aiohttp
|
|
from bot.utils.ai_client import get_client_timeout
|
|
|
|
connector = _get_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 e:
|
|
logger.error("Embedding request error: %s", e)
|
|
return None
|
|
|
|
|
|
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:
|
|
from bot.utils.ai_client import ask_ai_simple
|
|
|
|
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 e:
|
|
logger.error("Summary generation error: %s", e)
|
|
|
|
return None
|