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, }