Cognitive-rag / 04_vectoredb / README.md
README.md
Raw

Vector Storage & Retrieval (04_vectoredb)

Handles semantic search across documents, conversation memory, and related chunks with modular retrieval strategies.

Quick Start — For LLM / Generation Integration

Your teammate's main entry point is get_unified_context_for_llm(). It handles retrieval, reranking, and memory in one call.

from CognitiveRAG.04_vectoredb import get_unified_context_for_llm

Choose a strategy: "dense", "sparse", or "hybrid"

context = get_unified_context_for_llm(
    query_embedding=embedding_vector,   # dense vector from 03_embedding
    query_text="user question here",    # raw query string
    strategy="hybrid",                  # "dense" | "sparse" | "hybrid"
)

# Feed to LLM ↓
chunks_for_llm = context["unified_context"]

What it returns

Key Description
unified_context All chunks merged & sorted — feed this to the LLM
document_chunks Document-only chunks
memory_chunks Conversation memory chunks (boosted & filtered)
related_chunks Prev/next chunks (when include_related=True)
strategy Strategy used, e.g. "hybrid+rerank"
doc_count / memory_count Result counts

Optional parameters

Parameter Default Description
doc_top_k 30 (from config) How many document chunks to retrieve
memory_top_k 3 How many conversation memories to retrieve
include_related False Also fetch prev/next adjacent chunks
related_types ["prev", "next"] Which related chunk types to include
memory_boost 0.5 Score boost applied to memory results
memory_min_score 0.2 Min score to keep a memory result

[!NOTE] Reranking is automatic. When RERANKER_ENABLED=true (default), cross-encoder reranking runs
after retrieval and trims results to RERANKER_TOP_N (default 10). No extra code needed.

Strategy details

  • Dense — Semantic similarity via NVIDIA NIM embeddings. Needs query_embedding.
  • Sparse — BM25-style keyword retrieval via Pinecone sparse index. Needs query_text.
  • Hybrid — Dense + Sparse merged with Reciprocal Rank Fusion (RRF). Needs both.

Features

  • Multi-Namespace Retrieval: Query documents and conversation memory in parallel
  • Modular Strategies: Dense, Sparse (BM25-style), and Hybrid (RRF) retrieval
  • Cross-Encoder Reranking: Automatic reranking with cross-encoder/ms-marco-MiniLM-L6-v2
  • LangSmith Tracing: Automatic tracing of all retrieval operations
  • Relational Chunking: Fetch prev/next adjacent chunks for context continuity
  • JSON Output: Export retrieval results for downstream integration
  • NVIDIA NIM Embeddings: Uses nvidia/llama-3.2-nv-embedqa-1b-v2 (2048 dim)

Files

File Description
retriever_multi.py Main entry point — get_unified_context_for_llm(), retrieve(), strategy functions
pinecone_client.py Core client for upserting/querying (includes SparseClient)
reranker.py Cross-encoder reranking logic
config.py Configuration for API keys, indexes, reranker, and LangSmith
upsert_db.py Batch upload to dense/sparse indexes (`--mode dense
test_multi_retrieval.py Interactive test script with JSON output

Indexes

Index Type Metric Model
rag-hybrid-index-v2 Dense cosine nvidia/llama-3.2-nv-embedqa-1b-v2
rag-sparse-index-v1 Sparse dotproduct pinecone-sparse-english-v0

Configuration (config.py)

Variable Default Description
PINECONE_INDEX_NAME rag-hybrid-index-v2 Dense index
SPARSE_INDEX_NAME rag-sparse-index-v1 Sparse index
EMBEDDING_MODEL nvidia/llama-3.2-nv-embedqa-1b-v2 Dense embeddings
TOP_K 30 Results to retrieve
RERANKER_ENABLED true Enable cross-encoder reranking
RERANKER_MODEL cross-encoder/ms-marco-MiniLM-L6-v2 Reranker model
RERANKER_TOP_N 10 Passages returned after reranking
MEMORY_BOOST 0.5 Score boost for memories
LANGCHAIN_PROJECT cognitive-rag-retrieval LangSmith project

Lower-Level API

If you only need raw retrieval without memory/reranking, use retrieve():

from CognitiveRAG.04_vectoredb import retrieve, RetrievalStrategy

result = retrieve(
    query_embedding=vec,
    query_text="query",
    strategy=RetrievalStrategy.HYBRID,  # .DENSE | .SPARSE | .HYBRID
    top_k=30,
)
chunks = result["retrieved_chunks"]

Or call individual functions directly:

from CognitiveRAG.04_vectoredb import retrieve_dense, retrieve_sparse, retrieve_hybrid

result = retrieve_hybrid(
    query_embedding=embedding,
    query_text="patient medications",
    top_k=30,
    rrf_k=60,
)

Upserting Documents

# Dense index (default)
python upsert_db.py

# Sparse index
python upsert_db.py --mode sparse --file "path/to/chunks.jsonl"

# Both indexes
python upsert_db.py --mode both

LangSmith Tracing

All retrieval functions are automatically traced. Traces appear in your LangSmith dashboard under project cognitive-rag-retrieval.

Functions traced:

  • get_unified_context — Main retrieval with memories
  • retrieve_dense / retrieve_sparse / retrieve_hybrid

JSON Output Format

{
  "query": "What meds is the patient on?",
  "retrieved_chunks": ["chunk text 1", "chunk text 2"],
  "chunk_metadata": [
    {"chunk_id": "...", "doc_id": "...", "score": 0.92, "source": "hybrid"}
  ],
  "memory_chunks": ["previous Q/A text"],
  "stats": {
    "document_count": 5,
    "memory_count": 2,
    "related_count": 10
  }
}

Environment Files

File Contents
.envpinecone Pinecone API key
.envnvidia NVIDIA NIM API key
.envlangsmith LangSmith API key
Powered by GitFront