EAP – Document Indexing

# ── STEP 3 + 4: Embed + index ─────────────────────────
        collection_id = f"col-{tenant_id}"
        logger.info("📦 Indexing Into Vector Search")
        t0 = perf_counter()

        index_result = await index_chunks(
            project_id=PROJECT_ID,
            location=LOCATION,
            collection_id=collection_id,
            chunks=all_chunks,
            event_status=event_status,
        )
        timings["indexing"] = perf_counter() - t0
        logger.info(f"✅ Indexing Complete [{_fmt(timings['indexing'])}]")
import logging
import hashlib
from typing import List, Dict

from rag_pipeline.indexing import VectorIndexClient, enforce_schema
from rag_pipeline.indexing.helpers import (
    build_data_schema,
    build_auto_embedding_vector_schema,
)

logger = logging.getLogger("VectorIndexingService")

# – --------------------------------------
# Schema Configuration
# – --------------------------------------

EMBEDDING_FIELD = "embedding"

DATA_SCHEMA = build_data_schema(
    properties={
        "file_name": {"type": "string"},
        "body": {"type": "string"},
    },
    required=["body"],
)

VECTOR_SCHEMA = build_auto_embedding_vector_schema(
    embedding_field=EMBEDDING_FIELD,
    text_template="{body}",
)

SCHEMA_PROPS = DATA_SCHEMA["properties"]
SCHEMA_REQ = set(DATA_SCHEMA.get("required", []))


# – --------------------------------------
# Helper: Deterministic ID (VERY IMPORTANT)
# – --------------------------------------

def generate_id(file_name: str, body: str) -> str:
    return hashlib.md5((file_name + body).encode()).hexdigest()


# – --------------------------------------
# Core Service Function (EVENT DRIVEN)
# – --------------------------------------

async def index_chunks(
    *,
    project_id: str,
    location: str,
    collection_id: str,
    chunks: List[Dict],
    event_status: str,   # 🔥 NEW
):
    """
    Event-driven indexing:
    - CREATE → insert (idempotent)
    - UPDATE → delete + insert
    - DELETE → delete all
    """

    client = VectorIndexClient(project_id=project_id, location=location)
    await client.initialize()

    try:
        logger.info(f"⚡ Event received: {event_status}")

        # – --------------------------------------
        # Ensure collection exists
        # – --------------------------------------
        try:
            await client.create_collection(
                collection_id=collection_id,
                data_schema=DATA_SCHEMA,
                vector_schema=VECTOR_SCHEMA,
            )
            logger.info(f"✅ Collection created: {collection_id}")

        except Exception as e:
            if "AlreadyExists" in type(e).__name__ or "409" in str(e) or "already exists" in str(e).lower():
                logger.info(f"ℹ️ Collection already exists, skipping: {collection_id}")
            else:
                raise

        # – --------------------------------------
        # DELETE EVENT
        # – --------------------------------------
        if event_status == "DELETE":
            logger.info(f"🗑 Deleting all data from collection: {collection_id}")

            deleted = await client.delete_all_data_objects(collection_id=collection_id)

            return {
                "collection_id": collection_id,
                "deleted_objects": deleted,
                "event_status": "DELETE",
            }

        # – --------------------------------------
        # PREPARE PAYLOAD
        # – --------------------------------------
        ingest_payload = []
        skipped = 0

        for chunk in chunks:
            item = {
                "file_name": chunk.get("file_name", ""),
                "body": chunk.get("body", ""),
            }

            item, missing, _ = enforce_schema(item, SCHEMA_PROPS, SCHEMA_REQ)

            if missing:
                skipped += 1
                continue

            # 🔥 FIX: deterministic ID
            item["id"] = generate_id(item["file_name"], item["body"])

            ingest_payload.append(item)

        logger.info(f"📦 Ingest ready: {len(ingest_payload)} ({skipped} skipped)")

        if not ingest_payload:
            raise ValueError("No valid chunks to ingest")

        # – --------------------------------------
        # UPDATE EVENT
        # – --------------------------------------
        if event_status == "UPDATE":
            logger.info(f"🔄 UPDATE → clearing old data for collection: {collection_id}")

            await client.delete_all_data_objects(collection_id=collection_id)

        # – --------------------------------------
        # CREATE / UPDATE → INGEST
        # – --------------------------------------
        try:
            await client.ingest_records(
                collection_id=collection_id,
                records=ingest_payload,
                id_field="id",
            )

        except Exception as e:
            if "AlreadyExists" in type(e).__name__ or "409" in str(e) or "already exists" in str(e).lower():
                logger.warning("⚠️ Duplicate data detected → skipping existing records")
            else:
                raise

        logger.info("✅ Vector operation completed successfully")

        return {
            "collection_id": collection_id,
            "ingested": len(ingest_payload),
            "skipped": skipped,
            "event_status": event_status,
        }

    finally:
        await client.close()

Leave a Reply

Your email address will not be published. Required fields are marked *