Home_assistant/backend/app/rag/ingest.py

from __future__ import annotations

import hashlib
from typing import Any

from qdrant_client.http import models as qm
from sqlalchemy import select
from sqlalchemy.orm import Session

from app.config import get_settings
from app.db.models import ChatSession, Document, DocumentChunk
from app.rag import embeddings
from app.rag.chunker import ChunkMetadata, chunk_blocks
from app.rag.extractors import extract_document
from app.rag.store import (
    COLLECTION_DOC_CHUNKS,
    COLLECTION_FACTS,
    COLLECTION_SUMMARIES,
    delete_by_filter,
    upsert_points,
)


async def index_memory_fact(
    *,
    fact_id: int,
    user_id: int,
    content: str,
    category: str,
    importance: int,
    active: bool = True,
) -> None:
    settings = get_settings()
    if not settings.rag_enabled or not active:
        return
    vectors = await embeddings.embed_texts([content])
    if not vectors:
        return
    upsert_points(
        COLLECTION_FACTS,
        [
            qm.PointStruct(
                id=int(fact_id),
                vector=vectors[0],
                payload={
                    "user_id": user_id,
                    "fact_id": fact_id,
                    "category": category,
                    "content": content,
                    "importance": importance,
                },
            )
        ],
    )


async def deactivate_memory_fact(fact_id: int) -> None:
    settings = get_settings()
    if not settings.rag_enabled:
        return
    delete_by_filter(
        COLLECTION_FACTS,
        [qm.FieldCondition(key="fact_id", match=qm.MatchValue(value=fact_id))],
    )


async def index_session_summary(session_id: int, summary: str) -> None:
    settings = get_settings()
    if not settings.rag_enabled or not summary.strip():
        return
    from app.db.base import SessionLocal

    user_id = 1
    db = SessionLocal()
    try:
        session = db.get(ChatSession, session_id)
        if session:
            user_id = session.user_id
    finally:
        db.close()

    vectors = await embeddings.embed_texts([summary])
    if not vectors:
        return
    upsert_points(
        COLLECTION_SUMMARIES,
        [
            qm.PointStruct(
                id=int(session_id),
                vector=vectors[0],
                payload={"user_id": user_id, "session_id": session_id, "summary": summary[:4000]},
            )
        ],
    )


async def ingest_document_file(
    db: Session,
    *,
    user_id: int,
    title: str,
    filename: str,
    raw_bytes: bytes,
) -> dict[str, Any]:
    settings = get_settings()
    extracted = extract_document(raw_bytes, filename)
    if not extracted.text:
        raise ValueError("Пустой документ")

    digest = hashlib.sha256(raw_bytes).hexdigest()
    chunks = chunk_blocks(extracted.blocks)
    if not chunks:
        raise ValueError("Не удалось извлечь текстовые фрагменты из документа")

    doc = Document(
        user_id=user_id,
        title=title or filename,
        filename=filename,
        mime_type=extracted.mime_type,
        content_hash=digest,
        size_bytes=len(raw_bytes),
        chunk_count=len(chunks),
        index_status="pending" if settings.rag_enabled else "stored",
    )
    db.add(doc)
    db.flush()

    chunk_rows: list[DocumentChunk] = []
    for idx, piece in enumerate(chunks):
        row = _chunk_row(
            doc,
            idx,
            piece,
            filename=filename,
            mime_type=extracted.mime_type,
        )
        db.add(row)
        chunk_rows.append(row)
    db.commit()
    db.refresh(doc)

    if settings.rag_enabled and chunks:
        await index_document_chunks(db, doc)
        doc.index_status = "indexed"
        db.commit()
    elif settings.rag_enabled:
        doc.index_status = "empty"
        db.commit()

    db.refresh(doc)
    return serialize_document(doc)


async def index_document_chunks(db: Session, doc: Document) -> None:
    rows = list(doc.chunks)
    if not rows:
        return
    vectors = await embeddings.embed_texts([row.content for row in rows])
    points: list[qm.PointStruct] = []
    for row, vector in zip(rows, vectors, strict=False):
        db.refresh(row)
        points.append(
            qm.PointStruct(
                id=int(row.id),
                vector=vector,
                payload=_chunk_payload(doc, row),
            )
        )
    upsert_points(COLLECTION_DOC_CHUNKS, points)


async def reindex_document(db: Session, doc: Document) -> dict[str, Any]:
    settings = get_settings()
    if not settings.rag_enabled:
        doc.index_status = "stored"
        db.commit()
        db.refresh(doc)
        return serialize_document(doc)

    delete_document_vectors(doc.id)
    doc.index_status = "pending"
    db.commit()
    await index_document_chunks(db, doc)
    doc.index_status = "indexed"
    doc.chunk_count = len(doc.chunks)
    db.commit()
    db.refresh(doc)
    return serialize_document(doc)


def delete_document_vectors(document_id: int) -> None:
    delete_by_filter(
        COLLECTION_DOC_CHUNKS,
        [qm.FieldCondition(key="document_id", match=qm.MatchValue(value=document_id))],
    )


def serialize_document(doc: Document) -> dict[str, Any]:
    return {
        "id": doc.id,
        "title": doc.title,
        "filename": doc.filename,
        "mime_type": doc.mime_type,
        "chunk_count": doc.chunk_count,
        "index_status": doc.index_status,
        "size_bytes": doc.size_bytes,
        "created_at": doc.created_at.isoformat() if doc.created_at else None,
    }


def _chunk_row(
    doc: Document,
    chunk_index: int,
    chunk: ChunkMetadata,
    *,
    filename: str,
    mime_type: str,
) -> DocumentChunk:
    return DocumentChunk(
        document_id=doc.id,
        chunk_index=chunk_index,
        content=chunk.content,
        source_filename=filename,
        mime_type=mime_type,
        page_start=chunk.page_start,
        page_end=chunk.page_end,
        section_title=chunk.section_title,
        section_path=chunk.section_path,
        paragraph_index=chunk.paragraph_index,
        source_ref=chunk.source_ref,
    )


def _chunk_payload(doc: Document, row: DocumentChunk) -> dict[str, Any]:
    return {
        "user_id": doc.user_id,
        "document_id": doc.id,
        "chunk_id": row.id,
        "chunk_index": row.chunk_index,
        "title": doc.title,
        "filename": doc.filename,
        "mime_type": row.mime_type or doc.mime_type,
        "source_filename": row.source_filename or doc.filename,
        "page_start": row.page_start,
        "page_end": row.page_end,
        "section_title": row.section_title,
        "section_path": row.section_path,
        "paragraph_index": row.paragraph_index,
        "source_ref": row.source_ref,
        "content": row.content,
    }