437 lines
16 KiB
Python
437 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Гибридный поиск: Meilisearch + ChromaDB + LLM суммаризация.
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import requests as _requests
|
||
|
||
# Сессия без прокси для внутренних сервисов (Meilisearch/ChromaDB на localhost)
|
||
_internal_session = _requests.Session()
|
||
_internal_session.trust_env = False
|
||
|
||
# Загружаем .env
|
||
def load_env():
|
||
env_file = Path.home() / ".openclaw" / ".env"
|
||
if env_file.exists():
|
||
with open(env_file) as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line and not line.startswith("#") and "=" in line:
|
||
key, _, val = line.partition("=")
|
||
os.environ.setdefault(key.strip(), val.strip())
|
||
|
||
load_env()
|
||
|
||
MEILI_URL = "http://127.0.0.1:7700"
|
||
CHROMA_PATH = str(Path(__file__).parent.parent / "data" / "chromadb")
|
||
COLLECTION_NAME = "snowbike_embeddings"
|
||
INDEX_NAME = "snowbike_messages"
|
||
MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
|
||
OPENROUTER_KEY = os.environ.get("OPENROUTER_API_KEY", "")
|
||
LLM_MODEL = "anthropic/claude-sonnet-4-5"
|
||
|
||
# Порог релевантности для ChromaDB (меньше = лучше совпадение)
|
||
CHROMA_DISTANCE_THRESHOLD = 1.2
|
||
|
||
# === СЕЗОННЫЙ ФИЛЬТР ===
|
||
# Маппинг ключевых слов → месяцы (включая соседние для контекста)
|
||
SEASON_KEYWORDS = {
|
||
# Месяцы
|
||
"январ": [12, 1, 2],
|
||
"феврал": [1, 2, 3],
|
||
"март": [2, 3, 4],
|
||
"апрел": [3, 4, 5],
|
||
"май": [4, 5, 6],
|
||
"июн": [5, 6, 7],
|
||
"июл": [6, 7, 8],
|
||
"август": [7, 8, 9],
|
||
"сентябр": [8, 9, 10],
|
||
"октябр": [9, 10, 11],
|
||
"ноябр": [10, 11, 12],
|
||
"декабр": [11, 12, 1],
|
||
# Сезоны
|
||
"зим": [12, 1, 2],
|
||
"весн": [3, 4, 5],
|
||
"лет": [6, 7, 8],
|
||
"осен": [9, 10, 11],
|
||
# Сезон катания
|
||
"сезон": [11, 12, 1, 2, 3, 4], # типичный сноубайк-сезон
|
||
"межсезон": [5, 6, 7, 8, 9, 10],
|
||
}
|
||
|
||
|
||
def detect_season_months(query: str) -> Optional[list]:
|
||
"""
|
||
Определяет месяцы из запроса.
|
||
Возвращает список номеров месяцев (1-12) или None.
|
||
"""
|
||
query_lower = query.lower()
|
||
for keyword, months in SEASON_KEYWORDS.items():
|
||
if keyword in query_lower:
|
||
return months
|
||
return None
|
||
|
||
|
||
SYSTEM_PROMPT = """Ты — помощник по сноубайкам. Отвечаешь на вопросы на основе сообщений из Telegram-группы «Сноубайк Россия».
|
||
|
||
Правила:
|
||
1. Отвечай только на русском языке
|
||
2. Используй только информацию из предоставленных сообщений
|
||
3. Если информации недостаточно — честно скажи об этом, перечисли что было упомянуто
|
||
4. Всегда указывай конкретику: ДАТЫ (в формате «в апреле 2024», «в январе 2023» и т.д.), имена (если есть), цифры
|
||
5. Агрегируй мнения: если несколько человек говорят одно и то же — обобщи
|
||
6. Будь конкретен: перечисли проблемы/решения по пунктам
|
||
7. При ответе на сезонные вопросы — обязательно указывай из каких месяцев/лет взяты данные"""
|
||
|
||
USER_PROMPT_TEMPLATE = """Вопрос: {question}
|
||
|
||
Найденные сообщения из чата «Сноубайк Россия»:
|
||
{context}
|
||
|
||
Дай связный ответ на вопрос. Если нашлись конкретные факты — перечисли их.
|
||
Обязательно упоминай даты сообщений (месяц и год) когда это важно для контекста."""
|
||
|
||
|
||
# Глобальные клиенты (инициализируются лениво)
|
||
_chroma_collection = None
|
||
_embed_model = None
|
||
|
||
|
||
def get_chroma_collection():
|
||
global _chroma_collection
|
||
if _chroma_collection is None:
|
||
import chromadb
|
||
client = chromadb.PersistentClient(path=CHROMA_PATH)
|
||
_chroma_collection = client.get_collection(COLLECTION_NAME)
|
||
return _chroma_collection
|
||
|
||
|
||
def get_embed_model():
|
||
global _embed_model
|
||
if _embed_model is None:
|
||
from sentence_transformers import SentenceTransformer
|
||
_embed_model = SentenceTransformer(MODEL_NAME)
|
||
return _embed_model
|
||
|
||
|
||
def search_meilisearch(query: str, topic_ids: Optional[list] = None, limit: int = 20,
|
||
season_months: Optional[list] = None) -> list:
|
||
"""Полнотекстовый поиск в Meilisearch через HTTP."""
|
||
try:
|
||
fetch_limit = limit * 6 if season_months else limit
|
||
params = {"q": query, "limit": fetch_limit}
|
||
|
||
if topic_ids:
|
||
filter_str = " OR ".join(f"topic_id = {t}" for t in topic_ids)
|
||
params["filter"] = filter_str
|
||
|
||
resp = _internal_session.post(
|
||
f"{MEILI_URL}/indexes/{INDEX_NAME}/search",
|
||
json=params,
|
||
timeout=10,
|
||
)
|
||
resp.raise_for_status()
|
||
hits = resp.json().get("hits", [])
|
||
|
||
if season_months:
|
||
# Разделяем на "в сезон" и "вне сезона", возвращаем сначала сезонные
|
||
in_season = []
|
||
out_season = []
|
||
for hit in hits:
|
||
date_str = hit.get("date", "")
|
||
try:
|
||
month = int(date_str[5:7]) if len(date_str) >= 7 else 0
|
||
except (ValueError, IndexError):
|
||
month = 0
|
||
if month in season_months:
|
||
in_season.append(hit)
|
||
else:
|
||
out_season.append(hit)
|
||
hits = (in_season + out_season)[:limit]
|
||
|
||
return hits[:limit]
|
||
except Exception as e:
|
||
print(f"Meilisearch ошибка: {e}", file=sys.stderr)
|
||
return []
|
||
|
||
|
||
def search_chromadb(query: str, topic_ids: Optional[list] = None, limit: int = 20,
|
||
season_months: Optional[list] = None) -> list:
|
||
"""Семантический поиск в ChromaDB с фильтрацией по дистанции."""
|
||
try:
|
||
model = get_embed_model()
|
||
collection = get_chroma_collection()
|
||
|
||
embedding = model.encode(query).tolist()
|
||
|
||
# Строим where-фильтр
|
||
conditions = []
|
||
|
||
if topic_ids:
|
||
if len(topic_ids) == 1:
|
||
conditions.append({"topic_id": topic_ids[0]})
|
||
else:
|
||
conditions.append({"$or": [{"topic_id": t} for t in topic_ids]})
|
||
|
||
if season_months:
|
||
# Используем поле month если доступно (после переиндексации)
|
||
# Для старых записей без month — пробуем с запасом и пост-фильтрацией
|
||
has_month_field = _check_month_field(collection)
|
||
if has_month_field:
|
||
if len(season_months) == 1:
|
||
conditions.append({"month": season_months[0]})
|
||
else:
|
||
conditions.append({"$or": [{"month": m} for m in season_months]})
|
||
|
||
where = None
|
||
if len(conditions) == 1:
|
||
where = conditions[0]
|
||
elif len(conditions) > 1:
|
||
where = {"$and": conditions}
|
||
|
||
fetch_n = limit if not (season_months and not _check_month_field(collection)) else limit * 4
|
||
|
||
kwargs = {
|
||
"query_embeddings": [embedding],
|
||
"n_results": min(fetch_n, collection.count()),
|
||
"include": ["documents", "metadatas", "distances"],
|
||
}
|
||
if where:
|
||
kwargs["where"] = where
|
||
|
||
result = collection.query(**kwargs)
|
||
|
||
hits = []
|
||
ids = result.get("ids", [[]])[0]
|
||
docs = result.get("documents", [[]])[0]
|
||
metas = result.get("metadatas", [[]])[0]
|
||
dists = result.get("distances", [[]])[0]
|
||
|
||
for msg_id, doc, meta, dist in zip(ids, docs, metas, dists):
|
||
if dist > CHROMA_DISTANCE_THRESHOLD:
|
||
continue
|
||
# Пост-фильтр по месяцу если нет поля month
|
||
if season_months and not _check_month_field(collection):
|
||
date_str = meta.get("date", "")
|
||
try:
|
||
month = int(date_str[5:7]) if len(date_str) >= 7 else 0
|
||
except (ValueError, IndexError):
|
||
month = 0
|
||
if month not in season_months:
|
||
continue
|
||
hits.append({
|
||
"id": int(msg_id),
|
||
"text": doc,
|
||
"topic_id": meta.get("topic_id"),
|
||
"topic_title": meta.get("topic_title", ""),
|
||
"date": meta.get("date", ""),
|
||
"from_id": meta.get("from_id", ""),
|
||
"_chroma_distance": dist,
|
||
})
|
||
|
||
return hits[:limit]
|
||
except Exception as e:
|
||
print(f"ChromaDB ошибка: {e}", file=sys.stderr)
|
||
return []
|
||
|
||
|
||
_month_field_cache = None
|
||
|
||
def _check_month_field(collection) -> bool:
|
||
"""Проверяем есть ли поле month в коллекции (кэшируем результат)."""
|
||
global _month_field_cache
|
||
if _month_field_cache is not None:
|
||
return _month_field_cache
|
||
try:
|
||
sample = collection.get(limit=1, include=["metadatas"])
|
||
metas = sample.get("metadatas", [])
|
||
if metas and "month" in metas[0]:
|
||
_month_field_cache = True
|
||
else:
|
||
_month_field_cache = False
|
||
except Exception:
|
||
_month_field_cache = False
|
||
return _month_field_cache
|
||
|
||
|
||
def merge_results(meili_hits: list, chroma_hits: list, limit: int = 15) -> list:
|
||
"""Объединяем и дедуплицируем результаты."""
|
||
seen_ids = set()
|
||
merged = []
|
||
|
||
# Meilisearch результаты — полнотекстовые, высокий приоритет
|
||
for hit in meili_hits:
|
||
msg_id = hit.get("id")
|
||
if msg_id and msg_id not in seen_ids:
|
||
hit["_source"] = "meilisearch"
|
||
merged.append(hit)
|
||
seen_ids.add(msg_id)
|
||
|
||
# ChromaDB результаты — семантические, дополняем
|
||
for hit in chroma_hits:
|
||
msg_id = hit.get("id")
|
||
if msg_id and msg_id not in seen_ids:
|
||
hit["_source"] = "chromadb"
|
||
merged.append(hit)
|
||
seen_ids.add(msg_id)
|
||
|
||
return merged[:limit]
|
||
|
||
|
||
def format_context(messages: list) -> str:
|
||
"""Форматируем контекст для LLM с явными датами."""
|
||
lines = []
|
||
for i, msg in enumerate(messages, 1):
|
||
date_full = msg.get("date", "")
|
||
# Форматируем дату как "15 апреля 2024"
|
||
date_display = _format_date_russian(date_full)
|
||
topic = msg.get("topic_title", "?")
|
||
author = msg.get("from_id", "?")
|
||
text = msg.get("text", "")
|
||
msg_id = msg.get("id", "?")
|
||
source = msg.get("_source", "")
|
||
|
||
lines.append(f"[{i}] ID:{msg_id} | Дата: {date_display} | Топик: {topic} | автор:{author} | источник:{source}")
|
||
lines.append(f" {text}")
|
||
lines.append("")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
MONTH_NAMES_RU = {
|
||
1: "января", 2: "февраля", 3: "марта", 4: "апреля",
|
||
5: "мая", 6: "июня", 7: "июля", 8: "августа",
|
||
9: "сентября", 10: "октября", 11: "ноября", 12: "декабря",
|
||
}
|
||
|
||
|
||
def _format_date_russian(date_str: str) -> str:
|
||
"""Преобразует ISO дату в читаемый русский формат."""
|
||
try:
|
||
# "2024-04-15T..."
|
||
year = int(date_str[:4])
|
||
month = int(date_str[5:7])
|
||
day = int(date_str[8:10])
|
||
return f"{day} {MONTH_NAMES_RU[month]} {year}"
|
||
except Exception:
|
||
return date_str[:10] if date_str else "?"
|
||
|
||
|
||
def ask_llm(question: str, context: str) -> str:
|
||
"""Задаём вопрос LLM через OpenRouter."""
|
||
if not OPENROUTER_KEY:
|
||
return "ОШИБКА: OPENROUTER_API_KEY не задан в ~/.openclaw/.env"
|
||
|
||
user_msg = USER_PROMPT_TEMPLATE.format(question=question, context=context)
|
||
|
||
response = _requests.post(
|
||
"https://openrouter.ai/api/v1/chat/completions",
|
||
headers={
|
||
"Authorization": f"Bearer {OPENROUTER_KEY}",
|
||
"Content-Type": "application/json",
|
||
"HTTP-Referer": "https://snowbike-rag.local",
|
||
"X-Title": "Snowbike RAG",
|
||
},
|
||
json={
|
||
"model": LLM_MODEL,
|
||
"messages": [
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": user_msg},
|
||
],
|
||
"max_tokens": 1500,
|
||
"temperature": 0.3,
|
||
},
|
||
timeout=30,
|
||
)
|
||
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
return data["choices"][0]["message"]["content"]
|
||
|
||
|
||
def search(query: str, topic_ids: Optional[list] = None, limit: int = 15) -> dict:
|
||
"""
|
||
Главная функция гибридного поиска.
|
||
|
||
Returns:
|
||
dict с полями: query, answer, sources, count, time_ms
|
||
"""
|
||
start_time = time.time()
|
||
|
||
# Определяем сезонные месяцы
|
||
season_months = detect_season_months(query)
|
||
if season_months:
|
||
print(f"Сезонный запрос, фильтр по месяцам: {season_months}", file=sys.stderr)
|
||
|
||
# Поиск
|
||
meili_hits = search_meilisearch(query, topic_ids, limit=20, season_months=season_months)
|
||
chroma_hits = search_chromadb(query, topic_ids, limit=20, season_months=season_months)
|
||
|
||
print(f"Meilisearch hits: {len(meili_hits)}, ChromaDB hits: {len(chroma_hits)}", file=sys.stderr)
|
||
|
||
merged = merge_results(meili_hits, chroma_hits, limit=limit)
|
||
|
||
if not merged:
|
||
return {
|
||
"query": query,
|
||
"answer": "К сожалению, по вашему запросу ничего не найдено в базе сообщений.",
|
||
"sources": [],
|
||
"count": 0,
|
||
"time_ms": int((time.time() - start_time) * 1000),
|
||
}
|
||
|
||
# Формируем контекст для LLM
|
||
context = format_context(merged)
|
||
|
||
# LLM ответ
|
||
answer = ask_llm(query, context)
|
||
|
||
# Источники
|
||
sources = [
|
||
{
|
||
"id": msg.get("id"),
|
||
"date": msg.get("date", "")[:10],
|
||
"date_display": _format_date_russian(msg.get("date", "")),
|
||
"topic": msg.get("topic_title", ""),
|
||
"author": str(msg.get("from_id", "")),
|
||
"text_preview": msg.get("text", "")[:120],
|
||
}
|
||
for msg in merged
|
||
]
|
||
|
||
elapsed_ms = int((time.time() - start_time) * 1000)
|
||
|
||
return {
|
||
"query": query,
|
||
"answer": answer,
|
||
"sources": sources,
|
||
"count": len(merged),
|
||
"time_ms": elapsed_ms,
|
||
"season_months": season_months,
|
||
}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
query = sys.argv[1] if len(sys.argv) > 1 else "проблемы ремень трикстер"
|
||
|
||
print(f"Запрос: {query}\n")
|
||
result = search(query)
|
||
|
||
print(f"Найдено источников: {result['count']}")
|
||
print(f"Время: {result['time_ms']} мс")
|
||
if result.get("season_months"):
|
||
print(f"Сезонный фильтр (месяцы): {result['season_months']}")
|
||
print("\n=== ОТВЕТ ===")
|
||
print(result["answer"])
|
||
print("\n=== ИСТОЧНИКИ ===")
|
||
for s in result["sources"][:5]:
|
||
print(f" [{s['id']}] {s['date_display']} | {s['topic']} | {s['text_preview'][:100]}...")
|