Deep Dive into LlamaIndex's RAG Pipeline and Pinecone Vector Database Integration

import os
import time
from typing import List, Dict, Any, Optional
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    Settings
)
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.embeddings.openai import OpenAIEmbedding
from pinecone import Pinecone, ServerlessSpec
import logging

# Configure logging for production debugging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ProductionRAGPipeline:
    """Full LlamaIndex + Pinecone RAG pipeline with error handling and observability."""

    def __init__(
        self,
        pinecone_api_key: str,
        pinecone_env: str,
        pinecone_index_name: str,
        openai_api_key: str,
        embed_model_name: str = "text-embedding-3-small",
        chunk_size: int = 512,
        chunk_overlap: int = 128
    ):
        # Validate required env vars
        if not all([pinecone_api_key, openai_api_key, pinecone_index_name]):
            raise ValueError("Missing required API keys or index name")

        # Set LlamaIndex global settings
        Settings.embed_model = OpenAIEmbedding(
            model=embed_model_name,
            api_key=openai_api_key,
            embed_batch_size=100  # Match Pinecone's max batch size
        )
        Settings.chunk_size = chunk_size
        Settings.chunk_overlap = chunk_overlap
        Settings.llm = None  # Disable LLM for ingestion-only pipeline

        # Initialize Pinecone client with retries
        self.pc = Pinecone(
            api_key=pinecone_api_key,
            environment=pinecone_env,
            # Add custom timeout for large batch operations
            timeout=30
        )

        # Create index if it doesn't exist (serverless spec for cost efficiency)
        if pinecone_index_name not in self.pc.list_indexes().names():
            logger.info(f"Creating new Pinecone index: {pinecone_index_name}")
            self.pc.create_index(
                name=pinecone_index_name,
                dimension=1536,  # Match OpenAI embedding dimension
                metric="cosine",
                spec=ServerlessSpec(
                    cloud="aws",
                    region="us-east-1"
                )
            )
            # Wait for index to initialize (max 60s)
            for _ in range(12):
                if self.pc.describe_index(pinecone_index_name).status['ready']:
                    break
                time.sleep(5)
            else:
                raise TimeoutError(f"Pinecone index {pinecone_index_name} failed to initialize")

        # Initialize Pinecone vector store
        self.pinecone_index = self.pc.Index(pinecone_index_name)
        self.vector_store = PineconeVectorStore(
            pinecone_index=self.pinecone_index,
            # Enable metadata filtering for time-bound queries
            metadata_fields=["source", "publish_date", "author"]
        )
        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)

        logger.info(f"RAG pipeline initialized with index: {pinecone_index_name}")

    def ingest_documents(self, docs_dir: str) -> int:
        """Ingest all documents from a directory, return number of nodes upserted."""
        try:
            # Load raw documents with error handling for corrupt files
            reader = SimpleDirectoryReader(
                docs_dir,
                required_exts=[".pdf", ".md", ".txt"],
                recursive=True,
                errors="ignore"  # Skip corrupt files instead of crashing
            )
            documents = reader.load_data()
            logger.info(f"Loaded {len(documents)} raw documents from {docs_dir}")

            # Create vector store index (triggers upsert to Pinecone)
            index = VectorStoreIndex.from_documents(
                documents,
                storage_context=self.storage_context,
                show_progress=True
            )

            # Return number of vectors upserted (from Pinecone stats)
            stats = self.pinecone_index.describe_index_stats()
            return stats.total_vector_count

        except Exception as e:
            logger.error(f"Document ingestion failed: {str(e)}", exc_info=True)
            raise

if __name__ == "__main__":
    # Load env vars (use dotenv in production)
    pipeline = ProductionRAGPipeline(
        pinecone_api_key=os.getenv("PINECONE_API_KEY"),
        pinecone_env=os.getenv("PINECONE_ENV", "us-east-1"),
        pinecone_index_name="llama-index-docs-prod",
        openai_api_key=os.getenv("OPENAI_API_KEY")
    )
    upserted = pipeline.ingest_documents("./data")
    print(f"Successfully upserted {upserted} vectors to Pinecone")

import os
import logging
import time
from datetime import datetime, timedelta
from typing import List, Optional
from llama_index.core import (
    VectorStoreIndex,
    StorageContext,
    Settings,
    QueryBundle
)
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAILLM
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import (
    SimilarityPostprocessor,
    MetadataPostprocessor
)
from pinecone import Pinecone

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RAGQueryEngine:
    """Production query engine with hybrid search, filters, and postprocessing."""

    def __init__(
        self,
        pinecone_api_key: str,
        pinecone_index_name: str,
        openai_api_key: str,
        llm_model: str = "gpt-4o",
        top_k: int = 5,
        similarity_cutoff: float = 0.7
    ):
        # Validate inputs
        if not all([pinecone_api_key, openai_api_key, pinecone_index_name]):
            raise ValueError("Missing required credentials or index name")

        # Configure LlamaIndex settings
        Settings.embed_model = OpenAIEmbedding(
            model="text-embedding-3-small",
            api_key=openai_api_key
        )
        Settings.llm = OpenAILLM(
            model=llm_model,
            api_key=openai_api_key,
            temperature=0.1,  # Low temperature for factual answers
            max_tokens=1024
        )

        # Initialize Pinecone and vector store
        pc = Pinecone(api_key=pinecone_api_key)
        pinecone_index = pc.Index(pinecone_index_name)
        self.vector_store = PineconeVectorStore(
            pinecone_index=pinecone_index,
            # Enable hybrid search (sparse + dense)
            hybrid_search=True,
            sparse_embed_model="bm25"  # Uses LlamaIndex's built-in BM25
        )
        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)

        # Initialize index from existing Pinecone data
        self.index = VectorStoreIndex.from_vector_store(
            vector_store=self.vector_store,
            storage_context=self.storage_context
        )

        # Configure retriever with metadata filters and hybrid search
        self.retriever = VectorIndexRetriever(
            index=self.index,
            top_k=top_k,
            # Default filter: only include documents from last 90 days
            filters={
                "publish_date": {
                    "$gte": (datetime.now() - timedelta(days=90)).strftime("%Y-%m-%d")
                }
            },
            # Enable hybrid search weighting (0.7 dense, 0.3 sparse)
            hybrid_weight=0.7
        )

        # Configure query engine with postprocessors
        self.query_engine = RetrieverQueryEngine(
            retriever=self.retriever,
            node_postprocessors=[
                # Filter out low-relevance nodes
                SimilarityPostprocessor(similarity_cutoff=similarity_cutoff),
                # Ensure only approved sources are used
                MetadataPostprocessor(
                    required_fields=["source"],
                    excluded_sources=["wikipedia.org"]  # Example exclusion
                )
            ]
        )
        logger.info(f"Query engine initialized for index: {pinecone_index_name}")

    def query(
        self,
        query_text: str,
        override_filters: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """Run a query with optional filter overrides, return answer and metadata."""
        try:
            # Update retriever filters if override provided
            if override_filters:
                self.retriever.filters = override_filters
                logger.info(f"Applied filter override: {override_filters}")

            # Execute query
            start_time = time.time()
            response = self.query_engine.query(query_text)
            latency_ms = (time.time() - start_time) * 1000

            # Extract source nodes and metadata
            source_nodes = []
            for node in response.source_nodes:
                source_nodes.append({
                    "text": node.node.get_content()[:200] + "...",  # Truncate for readability
                    "score": node.score,
                    "metadata": node.node.metadata
                })

            return {
                "answer": str(response),
                "latency_ms": round(latency_ms, 2),
                "source_nodes": source_nodes,
                "total_tokens": response.metadata.get("total_tokens", 0)
            }

        except Exception as e:
            logger.error(f"Query failed: {str(e)}", exc_info=True)
            raise

if __name__ == "__main__":
    engine = RAGQueryEngine(
        pinecone_api_key=os.getenv("PINECONE_API_KEY"),
        pinecone_index_name="llama-index-docs-prod",
        openai_api_key=os.getenv("OPENAI_API_KEY")
    )
    result = engine.query("What are the benefits of LlamaIndex's Pinecone integration?")
    print(f"Answer: {result['answer']}")
    print(f"Latency: {result['latency_ms']}ms")
    print(f"Sources: {len(result['source_nodes'])} nodes retrieved")

import os
import time
import logging
from typing import List, Dict, Any
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type
)
from pinecone import Pinecone, Index
from pinecone.exceptions import PineconeRateLimitError, PineconeApiException
from llama_index.core.schema import TextNode, BaseNode
from llama_index.vector_stores.pinecone import PineconeVectorStore

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CustomPineconeVectorStore(PineconeVectorStore):
    """Extended Pinecone vector store with custom batching and retry logic."""

    def __init__(
        self,
        pinecone_index: Index,
        batch_size: int = 100,
        max_retries: int = 5,
        metadata_fields: List[str] = None
    ):
        super().__init__(pinecone_index=pinecone_index)
        self.batch_size = batch_size
        self.max_retries = max_retries
        self.metadata_fields = metadata_fields or []

        # Configure tenacity retry for rate limits
        self._upsert_with_retry = retry(
            stop=stop_after_attempt(max_retries),
            wait=wait_exponential(multiplier=1, min=4, max=60),
            retry=retry_if_exception_type((PineconeRateLimitError, PineconeApiException)),
            after=lambda retry_state: logger.warning(
                f"Retrying upsert (attempt {retry_state.attempt_number})"
            )
        )(self._upsert_batch)

    def _upsert_batch(self, batch: List[Dict]) -> None:
        """Upsert a single batch of vectors to Pinecone."""
        try:
            self._pinecone_index.upsert(vectors=batch)
            logger.debug(f"Upserted batch of {len(batch)} vectors")
        except PineconeRateLimitError as e:
            logger.error(f"Rate limit hit: {str(e)}")
            raise
        except Exception as e:
            logger.error(f"Batch upsert failed: {str(e)}", exc_info=True)
            raise

    def add_nodes(self, nodes: List[BaseNode]) -> List[str]:
        """Override LlamaIndex's add_nodes to use custom batching and retries."""
        ids = []
        batch = []

        for node in nodes:
            # Extract metadata fields specified in init
            metadata = {
                k: v for k, v in node.metadata.items()
                if k in self.metadata_fields or not self.metadata_fields
            }
            # Add node text to metadata for retrieval
            metadata["text"] = node.get_content()[:1000]  # Truncate to avoid size limits

            # Prepare vector entry
            vector_entry = {
                "id": node.node_id,
                "values": node.embedding,
                "metadata": metadata
            }
            batch.append(vector_entry)
            ids.append(node.node_id)

            # Upsert when batch is full
            if len(batch) >= self.batch_size:
                try:
                    self._upsert_with_retry(batch)
                    batch = []
                except Exception as e:
                    logger.error(f"Failed to upsert batch after {self.max_retries} retries: {str(e)}")
                    raise

        # Upsert remaining batch
        if batch:
            try:
                self._upsert_with_retry(batch)
            except Exception as e:
                logger.error(f"Failed to upsert final batch: {str(e)}")
                raise

        logger.info(f"Added {len(ids)} nodes to Pinecone")
        return ids

if __name__ == "__main__":
    # Initialize custom vector store
    pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
    index = pc.Index("llama-index-docs-prod")

    custom_store = CustomPineconeVectorStore(
        pinecone_index=index,
        batch_size=100,
        max_retries=5,
        metadata_fields=["source", "publish_date", "author"]
    )

    # Example node creation (in practice, from LlamaIndex pipeline)
    test_node = TextNode(
        text="LlamaIndex integrates with Pinecone via the PineconeVectorStore class.",
        metadata={"source": "llama-index-docs", "publish_date": "2024-06-01"}
    )
    test_node.embedding = [0.1] * 1536  # Mock embedding
    custom_store.add_nodes([test_node])
    print("Custom upsert test passed")

from llama_index.vector_stores.pinecone import PineconeVectorStore

# Upsert to tenant-specific namespace
vector_store = PineconeVectorStore(
    pinecone_index=pinecone_index,
    namespace="acme-corp"  # All upserts go to this namespace
)

# Override namespace per query
query_engine = RetrieverQueryEngine.from_args(
    index,
    vector_store_kwargs={"namespace": "globex"}  # Query Globex's namespace
)

vector_store = PineconeVectorStore(
    pinecone_index=pinecone_index,
    hybrid_search=True,
    hybrid_weight=0.7  # 70% dense, 30% sparse
)

from llama_index.core.callbacks import CallbackManager, OpenTelemetryCallbackHandler
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider

trace.set_tracer_provider(TracerProvider())
callback_handler = OpenTelemetryCallbackHandler()
Settings.callback_manager = CallbackManager([callback_handler])