218 lines
8.6 KiB
Python
218 lines
8.6 KiB
Python
"""ذخیره‌ساز برداری 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.rows import dict_row
|
||
from psycopg.types.json import Jsonb
|
||
|
||
from .embeddings import get_embeddings
|
||
|
||
model = settings.embedding_model
|
||
wanted = {c["doc_id"]: c for c in chunks}
|
||
|
||
# row_factory لازم است: پایین‌تر row["model"]/row["doc_id"] خوانده می‌شود و
|
||
# پیش‌فرض psycopg3 تاپل برمی‌گرداند (خطای tuple indices must be integers).
|
||
with psycopg.connect(chatbot_dsn(), autocommit=True, connect_timeout=10,
|
||
row_factory=dict_row) 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
|
||
from psycopg.rows import dict_row
|
||
|
||
with psycopg.connect(chatbot_dsn(), autocommit=True, connect_timeout=10,
|
||
row_factory=dict_row) 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]
|