aramesh-backend/chatbot/vectorstore.py

"""ذخیره‌ساز برداری pgvector — جایگزین FAISS در مهاجرت به pgvector.

بردارها در یک دیتابیس جدا (aramesh_chatbot) روی همان نمونه Postgres سایت
نگهداری می‌شوند تا دیتابیس سایت فقط‌خواندنی بماند. همگام‌سازی افزایشی است:
فقط سطرهای جدید/تغییریافته امبد می‌شوند (کلید: doc_id پایدار + hash محتوا).
"""

import hashlib
import json
import logging
from typing import Any
from urllib.parse import urlsplit, urlunsplit

from .config import settings

log = logging.getLogger(__name__)


def _strip_driver(dsn: str) -> str:
    return dsn.replace("+psycopg", "").replace("+asyncpg", "")


def chatbot_dsn() -> str:
    """DSN دیتابیس اختصاصی چت‌بات — از DATABASE_URL با تعویض نام دیتابیس ساخته می‌شود."""
    parts = urlsplit(_strip_driver(settings.database_url))
    return urlunsplit(parts._replace(path=f"/{settings.chatbot_db_name}"))


def _admin_dsn() -> str:
    """DSN دیتابیس postgres برای ساخت دیتابیس چت‌بات (فقط یک‌بار)."""
    parts = urlsplit(_strip_driver(settings.database_url))
    return urlunsplit(parts._replace(path="/postgres"))


def ensure_schema() -> None:
    """ساخت دیتابیس و جدول بردارها — idempotent و در هر شروع اجرا می‌شود."""
    import psycopg
    from psycopg import errors as pg_errors

    try:
        with psycopg.connect(_admin_dsn(), autocommit=True, connect_timeout=10) as conn:
            try:
                conn.execute(f'CREATE DATABASE "{settings.chatbot_db_name}"')
                log.info("Created chatbot database '%s'", settings.chatbot_db_name)
            except pg_errors.DuplicateDatabase:
                pass
    except Exception as e:
        raise RuntimeError(f"Cannot reach Postgres to ensure chatbot database: {e}") from e

    with psycopg.connect(chatbot_dsn(), autocommit=True, connect_timeout=10) as conn:
        conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
        conn.execute(
            f"""
            CREATE TABLE IF NOT EXISTS chatbot_chunks (
                doc_id       TEXT PRIMARY KEY,
                content_hash TEXT NOT NULL,
                model        TEXT NOT NULL,
                content      JSONB NOT NULL,
                embedding    vector({settings.embedding_dim}) NOT NULL
            )
            """
        )
        conn.execute(
            """
            CREATE INDEX IF NOT EXISTS chatbot_chunks_embedding_hnsw
                ON chatbot_chunks USING hnsw (embedding vector_cosine_ops)
            """
        )


def sync_chunks(chunks: list[dict], force: bool = False) -> dict:
    """همگام‌سازی افزایشی قطعات با جدول pgvector.

    - سطرهایی که دیگر وجود ندارند حذف می‌شوند
    - سطرهای جدید یا تغییر یافته (hash یا مدل امبدینگ عوض شده) امبد و upsert می‌شوند
    - force=True همه سطرها را از نو امبد می‌کند
    خروجی: آمار عملیات برای لاگ.
    """
    import psycopg
    from psycopg.types.json import Jsonb

    from .embeddings import get_embeddings

    model = settings.embedding_model
    wanted = {c["doc_id"]: c for c in chunks}

    with psycopg.connect(chatbot_dsn(), autocommit=True, connect_timeout=10) as conn:
        existing: dict[str, str] = {}
        if not force:
            for row in conn.execute("SELECT doc_id, content_hash, model FROM chatbot_chunks").fetchall():
                if row["model"] == model:
                    existing[row["doc_id"]] = row["content_hash"]
        else:
            conn.execute("DELETE FROM chatbot_chunks")

        to_delete = [doc_id for doc_id in existing if doc_id not in wanted]
        if to_delete:
            conn.execute(
                "DELETE FROM chatbot_chunks WHERE doc_id = ANY(%s)",
                [to_delete],
            )

        to_upsert = [
            c for c in chunks
            if existing.get(c["doc_id"]) != c["content_hash"]
        ]
        inserted = updated = 0

        if to_upsert:
            vectors = get_embeddings().embed_documents([c["text"] for c in to_upsert])
            for chunk, vec in zip(to_upsert, vectors):
                literal = "[" + ",".join(f"{x:.7f}" for x in vec) + "]"
                conn.execute(
                    """
                    INSERT INTO chatbot_chunks (doc_id, content_hash, model, content, embedding)
                    VALUES (%s, %s, %s, %s, %s::vector)
                    ON CONFLICT (doc_id) DO UPDATE SET
                        content_hash = EXCLUDED.content_hash,
                        model = EXCLUDED.model,
                        content = EXCLUDED.content,
                        embedding = EXCLUDED.embedding
                    """,
                    (chunk["doc_id"], chunk["content_hash"], model, Jsonb(chunk), literal),
                )
                if chunk["doc_id"] in existing:
                    updated += 1
                else:
                    inserted += 1

    stats = {"inserted": inserted, "updated": updated, "deleted": len(to_delete),
             "unchanged": len(wanted) - len(to_upsert)}
    log.info("Vector sync: %s", stats)
    return stats


def fetch_all_chunks() -> list[dict]:
    """همه قطعات ذخیره‌شده (برای ساخت BM25 در حافظه)."""
    import psycopg

    with psycopg.connect(chatbot_dsn(), autocommit=True, connect_timeout=10) as conn:
        rows = conn.execute(
            "SELECT doc_id, content_hash, content FROM chatbot_chunks ORDER BY doc_id"
        ).fetchall()
    return [
        {
            "doc_id": r["doc_id"],
            "content_hash": r["content_hash"],
            **(r["content"] if isinstance(r["content"], dict) else json.loads(r["content"])),
        }
        for r in rows
    ]


def vector_count() -> int:
    """تعداد بردارهای ذخیره‌شده — برای healthz (خطا → ۰)."""
    try:
        import psycopg

        with psycopg.connect(chatbot_dsn(), autocommit=True, connect_timeout=5) as conn:
            return int(conn.execute("SELECT count(*) FROM chatbot_chunks").fetchone()[0])
    except Exception:
        return 0


def dense_search(query: str, top_k: int) -> list[tuple[dict, float]]:
    """جست‌وجوی معنایی SQL با فاصله کسینوسی — خروجی: (سند، شباهت).

    خطا (مثلاً دیتابیس هنوز آماده نیست) → فهرست خالی تا بازیابی فقط با BM25 ادامه یابد.
    """
    try:
        import psycopg
        from psycopg.rows import dict_row as _dict_row

        from .embeddings import get_embeddings

        vec = get_embeddings().embed_query(query)
        literal = "[" + ",".join(f"{x:.7f}" for x in vec) + "]"
        with psycopg.connect(chatbot_dsn(), row_factory=_dict_row, connect_timeout=10) as conn:
            rows = conn.execute(
                f"""
                SELECT doc_id, content, 1 - (embedding <=> %s::vector) AS sim
                FROM chatbot_chunks
                ORDER BY embedding <=> %s::vector
                LIMIT %s
                """,
                (literal, literal, top_k),
            ).fetchall()
        return [(_row_to_doc(r), float(r["sim"])) for r in rows]
    except Exception as e:
        log.warning("Dense search failed (%s) — falling back to keyword only", e)
        return []


def _row_to_doc(row: Any) -> dict:
    content = row["content"] if isinstance(row["content"], dict) else json.loads(row["content"])
    return {
        "doc_id": row["doc_id"],
        "text": content.get("text", ""),
        "type": content.get("type", ""),
        "title": content.get("title", ""),
        "url": content.get("url", ""),
    }


def chunk_content_hash(chunk: dict) -> str:
    """هش محتوای قطعه — تغییر متن یا متادیتا یعنی سطر باید به‌روز شود."""
    payload = json.dumps(
        {k: chunk[k] for k in ("text", "type", "title", "url")},
        ensure_ascii=False, sort_keys=True,
    )
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16]