296 lines
10 KiB
Python
296 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Гибридный поиск: Meilisearch + ChromaDB + LLM суммаризация.
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
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
|
||
|
||
SYSTEM_PROMPT = """Ты — помощник по сноубайкам. Отвечаешь на вопросы на основе сообщений из Telegram-группы «Сноубайк Россия».
|
||
|
||
Правила:
|
||
1. Отвечай только на русском языке
|
||
2. Используй только информацию из предоставленных сообщений
|
||
3. Если информации недостаточно — честно скажи об этом, перечисли что было упомянуто
|
||
4. Всегда указывай конкретику: даты, имена (если есть), цифры
|
||
5. Агрегируй мнения: если несколько человек говорят одно и то же — обобщи
|
||
6. Будь конкретен: перечисли проблемы/решения по пунктам"""
|
||
|
||
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) -> list:
|
||
"""Полнотекстовый поиск в Meilisearch через HTTP."""
|
||
try:
|
||
params = {"q": query, "limit": 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()
|
||
return resp.json().get("hits", [])
|
||
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) -> list:
|
||
"""Семантический поиск в ChromaDB с фильтрацией по дистанции."""
|
||
try:
|
||
model = get_embed_model()
|
||
collection = get_chroma_collection()
|
||
|
||
embedding = model.encode(query).tolist()
|
||
|
||
where = None
|
||
if topic_ids and len(topic_ids) == 1:
|
||
where = {"topic_id": topic_ids[0]}
|
||
elif topic_ids and len(topic_ids) > 1:
|
||
where = {"$or": [{"topic_id": t} for t in topic_ids]}
|
||
|
||
kwargs = {
|
||
"query_embeddings": [embedding],
|
||
"n_results": limit,
|
||
"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
|
||
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
|
||
except Exception as e:
|
||
print(f"ChromaDB ошибка: {e}", file=sys.stderr)
|
||
return []
|
||
|
||
|
||
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 = msg.get("date", "")[:10]
|
||
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} | Топик: {topic} | автор:{author} | источник:{source}")
|
||
lines.append(f" {text}")
|
||
lines.append("")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
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()
|
||
|
||
# Поиск
|
||
meili_hits = search_meilisearch(query, topic_ids, limit=20)
|
||
chroma_hits = search_chromadb(query, topic_ids, limit=20)
|
||
|
||
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],
|
||
"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,
|
||
}
|
||
|
||
|
||
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']} мс\n")
|
||
print("=== ОТВЕТ ===")
|
||
print(result["answer"])
|
||
print("\n=== ИСТОЧНИКИ ===")
|
||
for s in result["sources"][:5]:
|
||
print(f" [{s['id']}] {s['date']} | {s['topic']} | {s['text_preview'][:100]}...")
|