RAG (Retrieval-Augmented Generation) API¶

def index_documents(self, docs: Sequence[Document]) -> None:
    """Index a batch of Documents into the RAG knowledge base.

    The implementation is responsible for:
    - Filtering to text documents (skipping bytes content)
    - Chunking per document according to backend strategy
    - Computing embeddings for all chunks
    - Upserting into the vector store

    Args:
        docs: Sequence of Document instances to index

    Raises:
        ValueError: If documents are invalid or cannot be indexed
        RuntimeError: If embedding or storage operations fail

    """
    ...

def query(self, request: RAGQueryRequest) -> RAGQueryResponse:
    """Execute vector search in the knowledge base.

    Returns ranked hits pointing back to original documents.

    Args:
        request: Query parameters (text, top_k, filters)

    Returns:
        Response containing ranked RAGHit results

    Raises:
        ValueError: If query parameters are invalid
        RuntimeError: If search operation fails

    """
    ...

def index_documents(self, docs: Sequence[Document]) -> None:
    """Index a batch of Documents into the RAG knowledge base.

    Implementation:
        1. Convert Documents to chunks using ingestion module
        2. Compute embeddings for all chunk texts
        3. Atomic upsert into LanceDB (merge_insert with update/insert)

    Args:
        docs: Sequence of Document instances to index

    Raises:
        ValueError: If documents are invalid
        RuntimeError: If embedding or storage operations fail

    """
    # Convert documents to chunks
    chunks = chunks_from_documents(docs, indexable_types=self._indexable_types)

    if not chunks:
        logger.info("No chunks to index (empty or filtered documents)")
        return

    logger.info("Indexing %d chunks from %d documents", len(chunks), len(docs))

    # Extract texts for embedding
    texts = [c.text for c in chunks]

    # Compute embeddings with RETRIEVAL_DOCUMENT task type
    try:
        embeddings = self._embed_fn(texts, "RETRIEVAL_DOCUMENT")
    except Exception as e:
        msg = f"Failed to compute embeddings: {e}"
        raise RuntimeError(msg) from e

    if len(embeddings) != len(chunks):
        msg = f"Embedding count mismatch: got {len(embeddings)}, expected {len(chunks)}"
        raise RuntimeError(msg)

    # Prepare rows using Pydantic models (ensures schema consistency)
    rows: list[RagChunkModel] = []
    for chunk, emb in zip(chunks, embeddings, strict=True):
        rows.append(
            RagChunkModel(
                chunk_id=chunk.chunk_id,
                document_id=chunk.document_id,
                text=chunk.text,
                vector=np.asarray(emb, dtype=np.float32),
                metadata_json=json.dumps(chunk.metadata),
            )
        )

    # Atomic upsert using merge_insert (no race conditions)
    try:
        self._table.merge_insert(
            "chunk_id"
        ).when_matched_update_all().when_not_matched_insert_all().execute(rows)
        logger.info("Successfully indexed %d chunks (atomic upsert)", len(rows))
    except Exception as e:
        msg = f"Failed to upsert chunks to LanceDB: {e}"
        raise RuntimeError(msg) from e

def query(self, request: RAGQueryRequest) -> RAGQueryResponse:
    """Execute vector search in the knowledge base.

    Implementation:
        1. Embed the query text
        2. Run vector search in LanceDB
        3. Convert results to RAGHit objects

    Args:
        request: Query parameters (text, top_k, filters)

    Returns:
        Response containing ranked RAGHit results

    Raises:
        ValueError: If query parameters are invalid
        RuntimeError: If search operation fails

    """
    top_k = request.top_k

    # Embed query with RETRIEVAL_QUERY task type
    try:
        query_emb = self._embed_fn([request.text], "RETRIEVAL_QUERY")[0]
    except Exception as e:
        msg = f"Failed to embed query: {e}"
        raise RuntimeError(msg) from e

    query_vec = np.asarray(query_emb, dtype=np.float32)

    # Execute search using Arrow (zero-copy, no Pandas)
    try:
        q = self._table.search(query_vec).metric("cosine").limit(top_k)

        # Apply filters if provided
        # LanceDB supports SQL-like WHERE clauses for pre-filtering
        if request.filters:
            q = q.where(request.filters)

        # Execute and get results as Arrow table (zero-copy)
        arrow_table = q.to_arrow()
    except Exception as e:
        msg = f"LanceDB search failed: {e}"
        raise RuntimeError(msg) from e

    # Convert Arrow table to Python dicts (fast native method)
    # This is much faster than iterating over pandas rows
    hits: list[RAGHit] = []
    for row in arrow_table.to_pylist():
        # LanceDB exposes a distance column (usually "_distance")
        distance = float(row.get("_distance", 0.0))

        # Convert distance to similarity score
        # For cosine distance: distance ∈ [0, 2], similarity = 1 - distance ∈ [-1, 1]
        # Normalize to [0, 1] range: score = (1 - distance) / 2 + 0.5
        # Simplified: score = (2 - distance) / 2 = 1 - (distance / 2)
        # However, for typical use cases, distance should be in [0, 2] range
        # and we want higher similarity scores for lower distances
        score = 1.0 - distance

        # Extract and deserialize metadata
        metadata_json = row.get("metadata_json", "{}")
        try:
            meta = json.loads(metadata_json) if metadata_json else {}
        except json.JSONDecodeError:
            logger.warning("Failed to decode metadata JSON, using empty dict")
            meta = {}

        hits.append(
            RAGHit(
                document_id=row["document_id"],
                chunk_id=row["chunk_id"],
                text=row["text"],
                metadata=meta,
                score=score,
            )
        )

    logger.info("Found %d hits for query (top_k=%d)", len(hits), top_k)
    return RAGQueryResponse(hits=hits)

def get_all_post_vectors(self) -> tuple[list[str], np.ndarray]:
    """Retrieve IDs and Centroid Vectors for all indexed posts.

    Returns:
        (doc_ids, vectors_matrix)

    """
    # Fetch all vectors (Zero-copy Arrow)
    # In a real scenario, filter by "document_type" metadata if possible
    arrow_table = self._table.search().limit(None).to_arrow()

    doc_vectors: dict[str, list[np.ndarray]] = {}

    # Aggregate chunks by document ID
    for batch in arrow_table.to_batches():
        d = batch.to_pydict()
        ids = d["document_id"]
        vecs = d["vector"]
        # Assume we only process POSTs based on upstream logic or metadata checks

        for i, doc_id in enumerate(ids):
            if doc_id not in doc_vectors:
                doc_vectors[doc_id] = []
            doc_vectors[doc_id].append(vecs[i])

    if not doc_vectors:
        return [], np.array([])

    # Compute centroids (Mean of chunk vectors)
    final_ids = []
    final_vecs = []

    for doc_id, vec_list in doc_vectors.items():
        centroid = np.mean(np.stack(vec_list), axis=0)
        final_ids.append(doc_id)
        final_vecs.append(centroid)

    return final_ids, np.array(final_vecs)