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()