Init

wiki_embeddings.py

@@ -0,0 +1,171 @@
+"""
+Семантический поиск по Яндекс Вики через OpenAI embeddings + pgvector.
+"""
+import re
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+from openai import OpenAI
+
+load_dotenv(Path(__file__).parent / '.env')
+
+EMBED_MODEL = 'text-embedding-3-small'  # 1536 dims, быстро и дёшево
+
+_client = None
+
+
+def _openai() -> OpenAI:
+    global _client
+    if _client is None:
+        _client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
+    return _client
+
+
+def clean_wiki_markup(text: str) -> str:
+    """Убрать wiki-разметку, оставить чистый текст для embedding."""
+    if not text:
+        return ''
+    # {% ... %} блоки
+    text = re.sub(r'\{%[^%]*%\}', '', text)
+    # [[ссылки]] → оставить текст после |
+    text = re.sub(r'\[\[([^\]|]*)\|([^\]]*)\]\]', r'\2', text)
+    text = re.sub(r'\[\[([^\]]*)\]\]', r'\1', text)
+    # ((url текст)) → текст
+    text = re.sub(r'\(\(https?://\S+\s+([^)]+)\)\)', r'\1', text)
+    text = re.sub(r'\(\(https?://\S+\)\)', '', text)
+    # Markdown-разметка
+    text = re.sub(r'\*{1,3}([^*]+)\*{1,3}', r'\1', text)
+    text = re.sub(r'_{1,2}([^_]+)_{1,2}', r'\1', text)
+    # Таблицы и спец-символы wiki
+    text = re.sub(r'[#|]{2,}', '\n', text)
+    text = re.sub(r'^\s*[#>]+\s*', '', text, flags=re.MULTILINE)
+    # Лишние пробелы
+    text = re.sub(r'\n{3,}', '\n\n', text)
+    return text.strip()
+
+
+def embed_text(text: str) -> list[float]:
+    """Получить embedding для текста через OpenAI API."""
+    import openai as _openai_module
+    # Начинаем с 15000 символов, при ошибке обрезаем вдвое
+    limit = 15000
+    while limit >= 1000:
+        try:
+            resp = _openai().embeddings.create(model=EMBED_MODEL, input=text[:limit])
+            return resp.data[0].embedding
+        except _openai_module.BadRequestError as e:
+            if 'maximum context length' in str(e):
+                limit = limit // 2
+                continue
+            raise
+    raise ValueError(f'Не удалось уложиться в лимит токенов даже при 1000 символах')
+
+
+def embed_page(page: dict) -> list[float]:
+    """Сгенерировать embedding для страницы (title + content)."""
+    title = page.get('title', '')
+    content = clean_wiki_markup(page.get('content', '') or '')
+    combined = f'{title}\n\n{content}'
+    return embed_text(combined)
+
+
+def upsert_embeddings(db, pages: list[dict]) -> dict:
+    """
+    Сгенерировать и сохранить embeddings для новых/изменённых страниц.
+
+    Пропускает страницы, у которых content_hash не изменился.
+
+    Returns:
+        dict: {'embedded': int, 'skipped': int}
+    """
+    import json
+    import pendulum
+
+    if not pages:
+        return {'embedded': 0, 'skipped': 0}
+
+    # Получить текущие хеши из wiki_embeddings
+    db.cursor.execute('SELECT slug, content_hash FROM wiki_embeddings')
+    stored = {row[0]: row[1] for row in db.cursor.fetchall()}
+
+    embedded = 0
+    skipped = 0
+
+    for page in pages:
+        slug = page['slug']
+        new_hash = page['content_hash']
+        content = page.get('content', '') or ''
+
+        if not content.strip():
+            skipped += 1
+            continue
+
+        if stored.get(slug) == new_hash:
+            skipped += 1
+            continue
+
+        print(f'  ↑ embedding: {slug}')
+        content_text = clean_wiki_markup(content)
+        vector = embed_page({'title': page.get('title', ''), 'content': content})
+
+        db.cursor.execute("""
+            INSERT INTO wiki_embeddings
+                (pg_load_dttm, slug, title, content_text, content_hash, embedding)
+            VALUES (%s, %s, %s, %s, %s, %s)
+            ON CONFLICT (slug) DO UPDATE SET
+                pg_load_dttm = EXCLUDED.pg_load_dttm,
+                title        = EXCLUDED.title,
+                content_text = EXCLUDED.content_text,
+                content_hash = EXCLUDED.content_hash,
+                embedding    = EXCLUDED.embedding
+        """, (
+            pendulum.now('Europe/Moscow'),
+            slug,
+            page.get('title', ''),
+            content_text,
+            new_hash,
+            json.dumps(vector),
+        ))
+        db.conn.commit()
+        embedded += 1
+
+    return {'embedded': embedded, 'skipped': skipped}
+
+
+def search(db, query: str, limit: int = 5) -> list[dict]:
+    """
+    Семантический поиск по wiki_embeddings.
+
+    Args:
+        db: подключённый SupabaseManager
+        query: текстовый запрос на любом языке
+        limit: кол-во результатов
+
+    Returns:
+        Список {'slug', 'title', 'similarity', 'content_text'}
+    """
+    import json
+
+    query_vec = embed_text(query)
+
+    db.cursor.execute("""
+        SELECT
+            slug,
+            title,
+            content_text,
+            1 - (embedding <=> %s::vector) AS similarity
+        FROM wiki_embeddings
+        ORDER BY embedding <=> %s::vector
+        LIMIT %s
+    """, (json.dumps(query_vec), json.dumps(query_vec), limit))
+
+    rows = db.cursor.fetchall()
+    return [
+        {
+            'slug': row[0],
+            'title': row[1],
+            'content_text': row[2],
+            'similarity': round(float(row[3]), 4),
+        }
+        for row in rows
+    ]