Vector Store Usage Guide

Overview

This guide provides practical recipes and patterns for using the Vector Store module in real-world scenarios, such as implementing Retrieval-Augmented Generation (RAG), semantic search with filtering, and multi-tenant data isolation.

---

Recipe 1: Indexing Documents for RAG

This recipe demonstrates how to chunk a text document, generate embeddings using an embedding model, and store them with metadata for later retrieval.

import uuid
from typing import List

from nala.athomic.ai.embeddings import EmbeddingModelProtocol
from nala.athomic.database.factory import connection_manager_factory
from nala.athomic.database.vector.types import VectorRecord

async def index_content(
    content: str, 
    metadata: dict, 
    embedding_model: EmbeddingModelProtocol
):
    # 1. Get the Vector Store client
    store = connection_manager_factory.create().get_vector_store()

    # 2. Chunk the content (Simple overlap strategy)
    chunk_size = 1000
    overlap = 100
    chunks = [
        content[i : i + chunk_size] 
        for i in range(0, len(content), chunk_size - overlap)
    ]

    # 3. Generate Embeddings in Batch
    # It is much more efficient to embed multiple texts at once
    vectors = await embedding_model.embed_documents(chunks)

    # 4. Create Vector Records
    records = []
    for chunk, vector in zip(chunks, vectors):
        record = VectorRecord(
            id=str(uuid.uuid4()),
            vector=vector,
            payload={
                "content": chunk,
                "source": metadata.get("source"),
                "author": metadata.get("author"),
                "created_at": metadata.get("created_at")
            }
        )
        records.append(record)

    # 5. Upsert to the 'knowledge_base' collection
    # Ensure the collection exists first (idempotent operation)
    await store.create_collection("knowledge_base", dimension=1536)

    await store.upsert("knowledge_base", records)
    print(f"Indexed {len(records)} chunks.")

---

Recipe 2: Semantic Search with Metadata Filtering

This recipe shows how to perform a semantic search while restricting results to specific criteria (e.g., only documents from a specific author).

async def search_knowledge_base(query: str, author_filter: str = None):
    store = connection_manager_factory.create().get_vector_store()
    embedding_model = ... # Get your embedding model

    # 1. Convert the user query into a vector
    query_vector = await embedding_model.embed_query(query)

    # 2. Build filters (if applicable)
    # The dictionary format allows providers to translate to their native query language
    filters = {}
    if author_filter:
        filters["author"] = author_filter

    # 3. Execute Search
    results = await store.search(
        collection="knowledge_base",
        vector=query_vector,
        limit=5,
        score_threshold=0.8, # Only relevant matches
        filter=filters
    )

    # 4. Process Results
    for hit in results:
        print(f"Score: {hit.score:.4f} | Content: {hit.payload['content'][:100]}...")

---

Recipe 3: Implementing Multi-Tenancy

One of the strongest features of the Athomic Vector Store is its seamless support for multi-tenancy via the ContextAwareVectorStore wrapper.

Automatic Injection

If you have the context_aware wrapper enabled in your configuration, you do not need to pass the tenant ID manually. The system extracts it from the current execution context.

# In a FastAPI request handler or Task
from nala.athomic.context import context_vars

async def handle_request(user_query: str):
    # Assume context middleware has already set this
    # context_vars.set_tenant_id("tenant-123")

    store = connection_manager_factory.create().get_vector_store()

    # This search is AUTOMATICALLY scoped to "tenant-123".
    # Data from other tenants will strictly NOT be visible.
    results = await store.search("documents", vector=...)

Manual Override (Admin Mode)

If you are writing an admin tool or a background job that needs to access a specific tenant's data explicitly, you can bypass the context:

async def admin_audit_tenant(target_tenant_id: str):
    store = connection_manager_factory.create().get_vector_store()

    # Explicitly passing 'tenant' overrides the context variable
    results = await store.search(
        "documents", 
        vector=..., 
        tenant=target_tenant_id
    )

---

Best Practices

1. Collection Management

• Dimension Matching: Always ensure the dimension passed to create_collection matches exactly the output dimension of your embedding model (e.g., OpenAI text-embedding-3-small is 1536).
• Naming Convention: Use snake_case for collection names. Avoid using tenant IDs in collection names; use the tenant partitioning feature instead for better scalability.

2. Error Handling

The Vector Store methods can raise specific exceptions that you should handle: * CollectionNotFoundError: If you try to search in a collection that hasn't been created. * ServiceConnectionError: If the vector database is down.

from nala.athomic.services.exceptions import ServiceConnectionError

try:
    await store.search(...)
except ServiceConnectionError:
    # Implement fallback logic (e.g., basic keyword search in DB)
    return await keyword_search(...)

3. Batching Upserts

Avoid calling upsert in a loop for single records. Always batch your records into lists (e.g., 100-500 records per call) to reduce network overhead and improve throughput.