DEV Community: DEVunderdog

One Hot Encoding

DEVunderdog — Sun, 16 Nov 2025 07:45:38 +0000

ChatGPT describes the one-hot encoding as following:

One-Hot Encoding is a categorical feature transformation technique used to convert non-numerical labels (like “red”, “blue”, “green”) into numerical binary vectors, so machine-learning models can understand them.

Now the thing with this approach is that what if we have 200s of categorical values, you are not going to create individual 200 columns for each categorical values.

import pandas as pd
import numpy as np

data = pd.read_csv('test.csv', usecols=['X1', 'X2', 'X3', 'X4', 'X5','X6'])
data.head()

This is what this dataset consists

Now we don't know what those categories emphasize, and we also don't know how much of those individual unique categorical value exists so let's try to figure out.

for col in data.columns:
   print(col, ": ", len(data[col].unique()), ' labels)

So we got the count of how much of unique categorical values exists.

Now let's say we convert all those unique categorical values into individual columns then

pd.get_dummies(data, drop_first=True).shape

That would yield (4209, 121) which means would have 121 columns which increases the dimension and that's not what we want.

So let's find top 10 most frequent categories specifically for X2 columns

data.X2.value_counts().sort_values(ascending=False).head(10)

Let's store those each individual top categorical values into index.

top_10 = [x for x in data.X2.value_counts().sort_values(ascending=False).head(10).index]

That would yield
['as', 'ae', 'ai', 'm', 'ak', 'r', 'n', 's', 'f', 'e']

Now let's have binary value that justify existence of categorical values

for label in top_10:
  data[label] = np.where(data['X2']==label, 1, 0)

data[['X2']+top_10].head(20)

Let's finalize that into function

def one_hot_top(df, variable, top_x_labels):
  for label in top_x_labels:
    df[variable+'_'+label] = np.where(data[variable]==label, 1, 0)

data = pd.read_csv('test.csv', usecols=['X1', 'X2', 'X3', 'X4', 'X5','X6'])

one_hot_top(data, 'X2', top_10)
data.head()

So basically you can notice that the categories which we consider we create a individual columns for those, and justify their existence using 0s and 1s if all those columns are 0s then that's not category we care, but if its 1 then that means that's the category.

I have shared my learning here, open for feedback.

Thanks :)

Socials:

I have shared my learnings, hope it is helpful to you.

DEVunderdog — Sat, 23 Aug 2025 15:41:48 +0000

DEVunderdog

Aug 23 '25

How I Cut AWS RAG System Response Time from 30 Seconds to 100ms

#go #aws #rag #api

Comments

19 min read

How I Cut AWS RAG System Response Time from 30 Seconds to 100ms

DEVunderdog — Sat, 23 Aug 2025 15:10:27 +0000

A deep dive into transforming a painful user experience into seamless interactions using resource pooling and state machines

The Pain Point That Changed Everything

Picture this: You've just built what you think is an elegant RAG (Retrieval Augmented Generation) system using AWS Bedrock, OpenSearch Serverless, and S3. The architecture looks clean on paper. Then you hit the "create knowledge base" button and... wait. And wait. Thirty seconds later, your users are still staring at a loading spinner.

This isn't just poor UX—it's a business killer in today's instant-gratification world.

The Hidden AWS Reality

Here's what the AWS documentation doesn't make crystal clear:

OpenSearch Serverless indexes take 15-20 seconds to become operationally searchable after creation.

This single constraint cascaded into a nightmare of complexity:

Create OpenSearch index → 15-20 seconds of mandatory wait time
Create AWS Knowledge Base → Remote API call with potential failures
Create Data Source → Another remote API call
Configure S3 inclusion prefixes → State management complexity
Handle failures at any step → Rollback hell

Each step could fail independently, creating inconsistent states between our database and AWS resources. The traditional synchronous approach was doomed from the start.

The "Aha" Moment: Think Like a Restaurant

The breakthrough came from an unexpected analogy. High-end restaurants don't start cooking your meal when you order—they prep ingredients in advance. They maintain pools of ready-to-use components.

What if we pre-provisioned AWS resources and maintained them in pools?

The Solution: Resource Pooling with State Machines

I implemented a multi-layered pool system that transforms the user experience:

Before: Click → Wait 30+ seconds → Maybe get a knowledge base
After: Click → Instant assignment from pre-warmed pool → Seamless experience

Architecture Overview

┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐
│   User Layer    │    │  Business Logic │    │  Resource Pools │
│                 │    │                 │    │                 │
│ Instant Response│◄──►│ Assign from Pool│◄──►│ Pre-warmed AWS  │
│                 │    │                 │    │    Resources    │
└─────────────────┘    └─────────────────┘    └─────────────────┘
                                                       │
                                               ┌───────▼────────┐
                                               │ Background     │
                                               │ Reconciliation │
                                               │ Workers        │
                                               └────────────────┘

The Three-Layer Pool System

Layer 1: OpenSearch Index Pool

CREATE TABLE `opensearch_indexes` (
    `id` BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
    `vector_index` VARCHAR(500) NOT NULL,
    `vector_field` VARCHAR(255) NOT NULL,
    `status` ENUM('PROVISIONING', 'AVAILABLE', 'ASSIGNED', 'DESTROYED', 'CLEANUP', 'FAILED') NOT NULL,
    `usable_at` TIMESTAMP NULL -- Critical for 15-20s warming period
);

Layer 2: AWS Knowledge Base Pool

CREATE TABLE `aws_knowledge_bases` (
    `id` BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
    `aws_kb_name` VARCHAR(100) NOT NULL,
    `index_id` BIGINT UNSIGNED NOT NULL,
    `knowledge_base_resource_id` VARCHAR(255) NULL,
    `status` ENUM('PROVISIONING', 'AVAILABLE', 'ASSIGNED', 'DESTROYED', 'CLEANUP', 'FAILED') NOT NULL,
    FOREIGN KEY (`index_id`) REFERENCES `opensearch_indexes`(`id`)
);

Layer 3: User Assignment Layer

CREATE TABLE `user_knowledge_bases` (
    `id` BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
    `name` VARCHAR(255) NOT NULL,
    `user_id` INT UNSIGNED NOT NULL,
    `aws_kb_id` BIGINT UNSIGNED NOT NULL,
    FOREIGN KEY (`aws_kb_id`) REFERENCES `aws_knowledge_bases`(`id`)
);

The State Machine That Changes Everything

Each resource follows a deterministic state flow:

PROVISIONING → AVAILABLE → ASSIGNED → CLEANUP → DESTROYED
      ↓             ↑
    FAILED    (Reconciliation)

This simple state machine eliminates the complexity of managing partial states and failed operations.

Each resource follows a clear state machine:

PROVISIONING: Resource is being created
AVAILABLE: Ready to be assigned to users
ASSIGNED: Currently in use by a user
CLEANUP: Marked for deletion
DESTROYED: Deleted (soft delete)
FAILED: Provisioning failed, needs cleanup

This approach provides several key benefits:

Index Manager

Index Manager is provision manager which handles the provisioning of Opensearch Indexes in the collections.

Property and Instantiation

Let's create a IndexManager which describes the provision manager and its property and let's try to have constructor which is utility to instantiate it.

The awsClient is a custom dependency which is being injected for handling AWS Operations. The store dependency which handles database operations, folks who would be familiar with sqlc would have an easy day to understand about this.

type IndexManager struct {
    awsClient                *aws_cloud.AWSCloudClients
    store                    database.Store
    baseLogger               *logger.Logger
    minSize                  int
    maxSize                  int
    maxConcurrentProvisioner int

    reconcileTrigger chan struct{}
}

func NewIndexManager(
    awsClient *aws_cloud.AWSCloudClients,
    store database.Store,
    baseLogger *logger.Logger,
) *IndexManager {
    return &IndexManager{
        awsClient:                awsClient,
        store:                    store,
        baseLogger:               baseLogger,
        minSize:                  globals.MinIndexPoolSize,
        maxSize:                  globals.MaxIndexPoolSize,
        maxConcurrentProvisioner: globals.MaxConcurrentProvisioner,

        reconcileTrigger: make(chan struct{}, 1),
    }
}

Provision a new index

It is a method of IndexManager and it basically creates a new index in Opensearch.

func (m *IndexManager) provisionNewIndex(ctx context.Context) error {
    indexName := utils.GenerateUniqueIndexName(true)
    vectorField := globals.VectorFieldName

    result, err := m.store.CreateOpensearchIndex(ctx, database.CreateOpensearchIndexParams{
        VectorIndex: indexName,
        VectorField: vectorField,
    })
    if err != nil {
        return fmt.Errorf("error initiating opensearch indexes in database: %w", err)
    }

    lastId, err := result.LastInsertId()
    if err != nil {
        return fmt.Errorf("error fetching last inserted id in opensearch indexs: %w", err)
    }

    err = m.awsClient.CreateIndexAndWaitForReady(ctx, indexName, vectorField)
    if err != nil {
        updateErr := m.store.UpdateOpensearchIndexStatus(ctx, database.UpdateOpensearchIndexStatusParams{
            Status: database.OpensearchIndexesStatusFAILED,
            ID:     uint64(lastId),
        })
        if updateErr != nil {
            return fmt.Errorf("error updating opensearch index status to failed due to err (%w): %w", err, updateErr)
        }
        return fmt.Errorf("error creating opensearch index remotely: %w", err)
    }

    err = m.store.FinalizeOpensearchIndex(ctx, uint64(lastId))
    if err != nil {
        return fmt.Errorf("error finalizing opensearch indexes: %w", err)
    }

    return nil
}

GenerateUniqueIndexName is utility function that actually generates unique name for indexes.
CreateOpensearchIndex is database operation function, which basically creates an Opensearch Index in the database with PROVISIONING status.
CreateIndexAndWaitForReady is AWS operation function which actually creates an index in Opensearch and make sure it has been created.
- if the remote operation got failed, then we actually marked the state of which was created in initially in database to FAILED using function UpdateOpensearchIndexStatus
FinalizeOpensearchIndex function just marks the state of initially created index to AVAILABLE such that its available to be assigned and also fill up usable_at column that justify warming of Opensearch Indexes.

Reconciling Index Pool

func (m *IndexManager) reconcileIndexPool(ctx context.Context) (bool, bool, error) {
    readyCount, err := m.store.CountReadyAvailableIndexes(ctx)
    if err != nil {
        return false, false, fmt.Errorf("error counting indexes by status: %w", err)
    }

    warmingCount, err := m.store.CountWarmingAvailableIndexes(ctx)
    if err != nil {
        return false, false, fmt.Errorf("error counting warming indexes: %w", err)
    }

    provisioningCount, err := m.store.CountProvisioningIndexes(ctx)
    if err != nil {
        return false, false, fmt.Errorf("failed to count provisioning indexes: %w", err)
    }

    if readyCount >= int64(m.minSize) {
        m.baseLogger.Info().Msgf("Index pool is healthy: Ready=%d", readyCount)
        return true, false, nil
    }

    totalPotentialPool := readyCount + warmingCount + provisioningCount

    if totalPotentialPool >= int64(m.minSize) {
        return false, false, nil
    }

    needed := m.minSize - int(totalPotentialPool)

    m.baseLogger.Info().Msgf(
        "index pool state: Available=%d, Warming=%d, Provisioning=%d. Need to create %d",
        readyCount,
        warmingCount,
        provisioningCount,
        needed,
    )

    g, gCtx := errgroup.WithContext(ctx)

    g.SetLimit(m.maxConcurrentProvisioner)

    for i := 0; i < needed; i++ {
        g.Go(func() error {
            m.baseLogger.Info().Msg("dispatching index provisioner task")

            err := m.provisionNewIndex(gCtx)
            if err != nil {
                m.baseLogger.Error().Err(err).Msg("failed to provision new index")
                return err
            }
            m.baseLogger.Info().Msg("successfully provisioned index")
            return nil
        })
    }

    if err := g.Wait(); err != nil {
        return false, false, fmt.Errorf("reconcilation failed during provision of indexs: %w", err)
    }

    m.baseLogger.Info().Msg("index reconcilation cycle finished")

    finalAvailableCount, err := m.store.CountIndexesByStatus(ctx, database.OpensearchIndexesStatusAVAILABLE)
    if err != nil {
        return false, false, fmt.Errorf("error countintt available indexes after provisioning: %w", err)
    }

    return finalAvailableCount >= int64(m.minSize), true, nil
}

using database query we count indexes in 'AVAILABLE' status and aren't warming up. What I meant by warming up is that it usually takes 15 to 20 seconds for searchable, hence we consider usable_at column too for that, so for that query would be something like this (basically defined as method CountReadyAvailableIndex)
```
    -- name: CountReadyAvailableIndexes :one
    select count(*) from opensearch_indexes
    where status = 'AVAILABLE'
        and usable_at is not null
        and usable_at < now() - interval 1 minute;
```

using database query we also count warming available indexes using following query.

    -- name: CountWarmingAvailableIndexes :one
    select count(*) from opensearch_indexes
    where status = 'AVAILABLE'
        and usable_at is not null
        and usable_at >= now() - interval 1 minute;

using database query we also count provisioning indexes using following query.

    -- name: CountProvisioningIndexes :one
    select count(*) from opensearch_indexes
    where status = 'PROVISIONING' and created_at > now() - interval 10 minute;

We are calculating total pool of indexes by adding all the counts, we are kind of conscious about over-provisioning hence we took this step.
If calculated pool is greater than minimum size than we return.
We then calculated the needed number of indexes
To simply handle the concurrency in golang we are using errgroup because,
- if any goroutine in the group returns an error, the entire group is considered failed, and we can handle the error appropriately.
- if one goroutines fails, context get cancelled, signalling other goroutines to stop their work early.
- Wait() method block until all goroutines complete.
We create a grouping contexted using errgroup.WithContext(ctx) additionally we also limit the concurrency using SetLimit method.
We run loop until the needed mark is met, and we start goroutine for each count and each goroutine executed provisionNewIndex function.
And we wait for all goroutine to complete using Wait() method.
And then we calculate the final AVAILABLE count for indexes and we make sure its above the minimum pool size.

Trigger Reconcilation

We would be reconciling and maintaining the pool of indexes periodically, but there would be instances where we need to trigger sometimes, hence you notice in NewIndexManager we created a buffered channel of size 1 and its for better communication and triggering, and the size describes that its not for subsequent triggering but just we balance out with manual triggering and periodic reconciliation.
triggerReconcilation does that.

func (m *IndexManager) triggerReconcilation() {
    select {
    case m.reconcileTrigger <- struct{}{}:
        m.baseLogger.Info().Msg("index reconcilation successfully trigger")
    default:
        m.baseLogger.Info().Msg("index reconcilation trigger ignored, reconcilation is already running")
    }
}

Reconcilation Worker

func (m *IndexManager) reconcilationWorker(ctx context.Context) error {
    ticker := time.NewTicker(5 * time.Minute)
    defer ticker.Stop()

    if _, _, err := m.reconcileIndexPool(ctx); err != nil {
        m.baseLogger.Error().Err(err).Msg("initial reconcilation failed")
        return err
    }

    for {
        select {
        case <-ctx.Done():
            m.baseLogger.Info().Msg("index reconcilation worker shutting down")
            return ctx.Err()

        case <-ticker.C:
            m.baseLogger.Info().Msg("periodic trigger of reconcilation of indexes")
            if _, _, err := m.reconcileIndexPool(ctx); err != nil {
                m.baseLogger.Error().Err(err).Msg("periodic index reconcilation failed")
                return err
            }

        case <-m.reconcileTrigger:
            m.baseLogger.Info().Msg("event-driven index trigger: starting reconcilation")

            select {
            case <-m.reconcileTrigger:
            default:
            }

            if _, _, err := m.reconcileIndexPool(ctx); err != nil {
                m.baseLogger.Error().Err(err).Msg("event-driven index reconcilation failed")
            }
        }
    }
}

We have worker which periodically hit the function reconcileIndexPool so basically by listening on ticker channel.
Similarly for manual triggering we are listening on reconcileTrigger channel which we initialize earlier, listen on that.

Claiming Indexes

We have query which basically runs to claim indexes based on database state

-- name: ClaimAvailableIndex :one
select
    id, vector_index, vector_field
from opensearch_indexes
    where status = 'AVAILABLE'
    and usable_at is not null
    and usable_at < now() - interval 1 minute
order by rand()
limit 1
for update skip locked;

the reason we have kept for update skip locked is because we are choosing randomly from record we fetched based on conditions, if its locked we moved on to next record to select.
And once we fetch that record, we mark it as assigned within the transaction using following query.

-- name: MarkIndexAsAssigned :exec
update opensearch_indexes
set
    status = 'ASSIGNED', assigned_at = now()
where id = ?;

So this is whole transaction defined as TxClaimAvailableIndex

func (store *SqlStore) TxClaimAvailableIndex(
    ctx context.Context,
) (*ClaimAvailableIndexRow, error) {
    var response ClaimAvailableIndexRow
    err := store.execTx(ctx, func(q *Queries) error {
        index, err := q.ClaimAvailableIndex(ctx)
        if err != nil {
            return fmt.Errorf("failed to claim available index: %w", err)
        }

        err = q.MarkIndexAsAssigned(ctx, index.ID)
        if err != nil {
            return fmt.Errorf("failed to mark index as assigned: %w", err)
        }
        response = index

        return nil

    })

    if err != nil {
        return nil, err
    }

    return &response, nil

}

After claiming index from database we trigger manually reconciliation of indexes.

func (m *IndexManager) ClaimAvailableIndex(ctx context.Context) (*database.ClaimAvailableIndexRow, error) {
    index, err := m.store.TxClaimAvailableIndex(ctx)
    if err != nil {
        return nil, err
    }

    go m.triggerReconcilation()

    return index, nil
}

Cleaning up indexes

func (m *IndexManager) cleanupRemoteResources(ctx context.Context, indexName string) error {
    err := m.awsClient.DeleteIndex(ctx, indexName)
    if err != nil {
        return fmt.Errorf("error deleting opensearch index remotely: %w", err)
    }

    return nil
}

func (m *IndexManager) CleanupIndexes(ctx context.Context) error {
    indexesForCleanup, err := m.store.FetchIndexesForCleanup(ctx)
    if err != nil {
        return fmt.Errorf("error fetching indexes for clean up: %w", err)
    }

    if len(indexesForCleanup) == 0 {
        m.baseLogger.Info().Msg("no indexes to cleanup")
        return nil
    }

    m.baseLogger.Info().Msgf("found %d indexes to cleanup", len(indexesForCleanup))

    g, gCtx := errgroup.WithContext(ctx)
    g.SetLimit(m.maxConcurrentProvisioner)

    for _, idx := range indexesForCleanup {
        indexToCleanup := idx

        g.Go(func() error {
            if err := gCtx.Err(); err != nil {
                m.baseLogger.Warn().Msgf("cleanup cancelled for index %s due to earlier error", indexToCleanup.VectorField)
                return err
            }

            err := m.cleanupRemoteResources(gCtx, indexToCleanup.VectorIndex)
            if err != nil {
                m.baseLogger.Error().
                    Err(err).
                    Str("index", indexToCleanup.VectorIndex).
                    Msg("failed to cleanup remote resources for index")
            }

            _, err = m.store.DeleteOpensearchIndex(gCtx, indexToCleanup.ID)
            if err != nil {
                m.baseLogger.Error().
                    Err(err).
                    Uint64("id", indexToCleanup.ID).
                    Msg("failed to mark index as destroyed in database")
                return fmt.Errorf("failed to update database status: %w", err)
            }

            m.baseLogger.Info().
                Str("index", indexToCleanup.VectorIndex).
                Msg("successfully cleaned up index")

            return nil
        })
    }

    if err := g.Wait(); err != nil {
        m.baseLogger.Error().Err(err).Msg("cleanup cycle completed with one or more failures")
    }

    return nil
}

We are basically handling remote index cleanup in Opensearch using cleanupRemoteResources method.
In CleanupIndexes we first fetch the indexes for cleanup and query is like this:

-- name: FetchIndexesForCleanup :many
select
    oi.*
from opensearch_indexes oi
left join
    aws_knowledge_bases akb on oi.id = akb.index_id
where
    oi.status = 'FAILED'
    or (oi.status = 'PROVISIONING' and oi.created_at < now() - interval 10 minute)
    or (
        oi.status = 'CLEANUP'
        and (akb.id is null or akb.status = 'DESTROYED')
    );

- We are doing joins because we wanted to make sure that indexes which are set for `CLEANUP` do have destroyed knowledge base other wise we would have cascading inconsistency, of course it would not be enforced by database, but we need to take care of that as we are working with soft-deletes which is actually a logical layer of deletion.

After fetching we trying to concurrently trying to cleanup remotely and once cleaned up remotely we try to soft-delete in database using this query

-- name: DeleteOpensearchIndex :execresult
update opensearch_indexes
set
    status = 'DESTROYED',
    deleted_at = current_timestamp
where id = ?;

CRON based Cleanup Worker

According to our requirements we are never going to run this application multiple instances, hence a single instance would run and we would try to initially do vertical scaling, because we are in kind of beta mode so not expecting much traffic hence we are setting CRON job within application to schedule Cleanup, but it would be drawback and not suitable if we are ran multiple instances of this application because then coordination would be complex and that would be another day problem to solve for.

func (m *IndexManager) CleanupIndexWorker(ctx context.Context) {
    location, err := time.LoadLocation("America/Chicago")
    if err != nil {
        m.baseLogger.Error().Err(err).Msg("cannot load the timezone")
        return
    }

    c := cron.New(cron.WithLocation(location))

    schedule := "30 3 * * *"
    _, err = c.AddFunc(schedule, func() {
        m.baseLogger.Info().
            Str("schedule", schedule).
            Str("timezone", location.String()).
            Msg("triggered scheduled index cleanup job")

        jobCtx, cancel := context.WithTimeout(context.Background(), 2*time.Hour)
        defer cancel()

        if err := m.CleanupIndexes(jobCtx); err != nil {
            m.baseLogger.Error().Err(err).Msg("scheduled index cleanup job failed")
        } else {
            m.baseLogger.Info().Msg("scheduled index cleanup job completed successfully")
        }
    })

    if err != nil {
        m.baseLogger.Error().Err(err).Msg("fatal: could not schedule cleanup job")
        return
    }

    m.baseLogger.Info().
        Str("schedule", schedule).
        Str("timezone", location.String()).
        Msg("index cleanup job scheduled successfully")

    c.Start()

    <-ctx.Done()

    m.baseLogger.Info().Msg("shutting down scheduled index cleanup worker")

    stopCtx := c.Stop()
    <-stopCtx.Done()

    m.baseLogger.Info().Msg("index cleanup scheduler stopped gracefully")
}

Run Index Manager

We define method to run Index Manager wrapped in method Run

func (m *IndexManager) Run(ctx context.Context) error {
    m.baseLogger.Info().Msg("starting index manager workers")

    g, gCtx := errgroup.WithContext(ctx)

    g.Go(func() error {
        err := m.reconcilationWorker(ctx)
        return err
    })

    g.Go(func() error {
        m.CleanupIndexWorker(gCtx)
        return gCtx.Err()
    })

    m.baseLogger.Info().Msg("all index manager workers are running")
    err := g.Wait()
    if err != nil && !errors.Is(err, context.Canceled) {
        m.baseLogger.Error().Err(err).Msg("index manager stopped due to an error")
        return err
    }

    m.baseLogger.Info().Msg("index manager stopped gracefully")
    return nil
}

We start in background as goroutine the reconcilationWorker and CleanupIndexWorker.

Knowledge Base Manager

Let's define Knowledge base Manager

type KbManager struct {
    awsClient                *aws_cloud.AWSCloudClients
    store                    database.Store
    config                   *utils.Config
    baseLogger               *logger.Logger
    indexManager             *IndexManager
    minSize                  int
    maxSize                  int
    maxConcurrentProvisioner int

    reconcileTrigger chan struct{}
}

Initiating Knowledge base manager

func NewKbManager(
    awsClient *aws_cloud.AWSCloudClients,
    store database.Store,
    baseLogger *logger.Logger,
    config *utils.Config,
    indexManager *IndexManager,
) *KbManager {
    return &KbManager{
        awsClient:                awsClient,
        store:                    store,
        baseLogger:               baseLogger,
        config:                   config,
        minSize:                  globals.MinKbPoolSize,
        maxSize:                  globals.MaxKbPoolSize,
        maxConcurrentProvisioner: globals.MaxConcurrentProvisioner,
        indexManager:             indexManager,

        reconcileTrigger: make(chan struct{}, 1),
    }
}

Provision New Knowledge Base

The reason we are injecting index manager into knowledge base manager because we need to claim indexes for knowledge bases

func (m *KbManager) provisionNewKnowledgeBase(ctx context.Context) error {
    m.baseLogger.Info().Msg("attempting to claim an available index for new knowledge base")
    claimedIndex, err := m.indexManager.ClaimAvailableIndex(ctx)
    if err != nil {
        if database.IsRecordNotFound(err) {
            m.baseLogger.Warn().Msg("no available index to claim")
            go m.indexManager.triggerReconcilation()
            return fmt.Errorf("no available index in the pool")
        }
        return fmt.Errorf("failed to claim available index: %w", err)
    }

    m.baseLogger.Info().Msg("successfully claimed indexed")

    dataSourceName := utils.GenerateUniqueIndexName(false)
    storagePrefix := utils.GenerateS3Prefix()
    kbName := utils.GenerateKnowledgeBaseName()

    inclusionPrefix := fmt.Sprintf("%s/%s/", m.config.AWSParentFolder, storagePrefix)

    result, err := m.store.CreateAwsKnowledgeBase(ctx, database.CreateAwsKnowledgeBaseParams{
        AwsKbName:      kbName,
        IndexID:        claimedIndex.ID,
        DataSourceName: dataSourceName,
        StoragePrefix:  inclusionPrefix,
    })
    if err != nil {
        return fmt.Errorf("error initiating knowledge base in database: %w", err)
    }

    lastId, err := result.LastInsertId()
    if err != nil {
        return fmt.Errorf("error fetching last inserted id in the knowledge base: %w", err)
    }

    kbResult, err := m.awsClient.CreateKnowledgeBase(ctx, aws_cloud.CreateKnowledgeBaseParams{
        Name:              &kbName,
        RoleARN:           &m.config.AWSIamRoleARN,
        S3Location:        &m.config.AWSDefaultS3Uri,
        EmbeddingModelARN: &m.config.AWSEmbeddingModelARN,
        CollectionARN:     &m.config.AWSCollectionARN,
        VectorIndexName:   &claimedIndex.VectorIndex,
        VectorField:       &claimedIndex.VectorField,
    })
    if err != nil {
        indexErr := m.store.MarkOpensearchIndexForCleanup(ctx, claimedIndex.ID)
        if indexErr != nil {
            return fmt.Errorf("error marking index for cleanup while rollbacking failure in knowledge base creation (%w): %w", err, indexErr)
        }

        _, rollbackErr := m.store.DeleteAwsKnowledgeBase(ctx, uint64(lastId))
        if rollbackErr != nil {
            return fmt.Errorf("error rollbacking knowledge base due to error(%w): %w", err, rollbackErr)
        }

        return fmt.Errorf("failed to create knowledge base: %w", err)
    }

    dataSourceResult, err := m.awsClient.CreateDataSource(
        ctx,
        dataSourceName,
        kbResult.KnowledgeBase.KnowledgeBaseId,
        &m.config.AWSBucketARN,
        &m.config.AWSBucketOwnerID,
        &inclusionPrefix,
    )
    if err != nil {
        _, kbErr := m.awsClient.DeleteKnowledgeBase(ctx, *kbResult.KnowledgeBase.KnowledgeBaseId)
        if kbErr != nil {
            return fmt.Errorf("error deleting knowledge base (due to - %w): %w", err, kbErr)
        }

        indexErr := m.store.MarkOpensearchIndexForCleanup(ctx, claimedIndex.ID)
        if indexErr != nil {
            return fmt.Errorf("error marking index for cleanup while rollbacking failure in knowledge base creation (%w): %w", err, indexErr)
        }

        _, rollbackErr := m.store.DeleteAwsKnowledgeBase(ctx, uint64(lastId))
        if rollbackErr != nil {
            return fmt.Errorf("error rollbacking knowledge base due to error(%w): %w", err, rollbackErr)
        }
        return fmt.Errorf("failed to create data source: %w", err)
    }

    err = m.store.FinalizedAwsKnowledgeBase(ctx, database.FinalizedAwsKnowledgeBaseParams{
        DataSourceID: sql.NullString{
            Valid:  true,
            String: *dataSourceResult.DataSource.DataSourceId,
        },
        KnowledgeBaseResourceID: sql.NullString{
            Valid:  true,
            String: *kbResult.KnowledgeBase.KnowledgeBaseId,
        },
        ID: uint64(lastId),
    })
    if err != nil {
        dataSourceErr := m.awsClient.DeleteDataSource(ctx, kbResult.KnowledgeBase.KnowledgeBaseId, dataSourceResult.DataSource.DataSourceId)
        if dataSourceErr != nil {
            return fmt.Errorf("error deleting data source (due to: %w): %w", err, dataSourceErr)
        }

        _, knowledgeBaseErr := m.awsClient.DeleteKnowledgeBase(ctx, *kbResult.KnowledgeBase.KnowledgeBaseId)
        if knowledgeBaseErr != nil {
            return fmt.Errorf("error deleting knowledge base (due to: %w): %w", err, knowledgeBaseErr)
        }

        indexErr := m.store.MarkOpensearchIndexForCleanup(ctx, claimedIndex.ID)
        if indexErr != nil {
            return fmt.Errorf("error marking index for cleanup while rollbacking failure in knowledge base creation (%w): %w", err, indexErr)
        }

        return fmt.Errorf("error finalizing knowledge base: %w", err)
    }

    return nil
}

Initially we are claiming indexes from index manager
Then we create AWS Knowledge base in our database using following query,

insert into aws_knowledge_bases (
    aws_kb_name,
    index_id,
    data_source_name,
    storage_prefix,
    status
) values (
    ?, ?, ?, ?, 'PROVISIONING'
);

The we create knowledge base actually in AWS, and if something fails we handle rollback by marking index for cleanup and delete the initiated state in database.
Then after creating knowledge base we need to create data source and then we finalize the knowledge base creation.

-- name: FinalizedAwsKnowledgeBase :exec
update aws_knowledge_bases
set
    data_source_id = ?,
    knowledge_base_resource_id = ?,
    status = 'AVAILABLE'
where id = ?;

again we handle rollback if finalization failed by deleting remote data sources, knowledge bases and marking opensearch indexes for cleanup.

Reconciling Knowledge Base Pool

Similar to what we do in index manager for reconciliation of pool we do the same for knowledge base
The only difference is we don't have warming instances for this.

func (m *KbManager) reconcileKnowledgeBasePool(ctx context.Context) error {
    availableCount, err := m.store.CountKnowledgeBaseByStatus(ctx, database.AwsKnowledgeBasesStatusAVAILABLE)
    if err != nil {
        return fmt.Errorf("failed to count available knowledge bases: %w", err)
    }

    provisioningCount, err := m.store.CountProvisioningKnowledgeBase(ctx)
    if err != nil {
        return fmt.Errorf("failed to count provisioning knowledge base: %w", err)
    }

    count := availableCount + provisioningCount

    if count >= int64(m.minSize) {
        return nil
    }

    needed := m.minSize - int(count)

    m.baseLogger.Info().Msgf(
        "knowledge base pool state: Available=%d, Provisioning=%d. Need to create %d",
        availableCount,
        provisioningCount,
        needed,
    )

    g, gCtx := errgroup.WithContext(ctx)

    g.SetLimit(m.maxConcurrentProvisioner)

    for i := 0; i < needed; i++ {
        g.Go(func() error {
            m.baseLogger.Info().Msg("dispatching knowledge base provisioner task")

            err := m.provisionNewKnowledgeBase(gCtx)
            if err != nil {
                m.baseLogger.Error().Err(err).Msg("failed to provision knowledge base")
                return err
            }
            m.baseLogger.Info().Msg("successfully provisioned a knowledge base")
            return nil
        })
    }

    if err := g.Wait(); err != nil {
        return fmt.Errorf("reconcilation failed during provisioning knowledge base: %w", err)
    }

    m.baseLogger.Info().Msg("knowledge base reconcilation cycle finished")

    return nil
}

Then whenever whatever components where knowledge base manager will be utilize they would be able to get available knowledge base using following method.

func (m *KbManager) AssignedAvailableKnowledgeBase(
    ctx context.Context,
    args database.TxAssignKnowledgeBaseParams,
) (*int64, error) {
    knowledgeBaseId, err := m.store.TxAssignKnowledgeBase(ctx, args)
    if err != nil {
        return nil, err
    }

    go m.triggerReconcilation()

    return knowledgeBaseId, nil
}

And we are triggering for reconciliation similar to index manager

func (m *KbManager) triggerReconcilation() {
    select {
    case m.reconcileTrigger <- struct{}{}:
        m.baseLogger.Info().Msg("knowledge base reconcilation successfully triggered")
    default:
        m.baseLogger.Info().Msg("knowledge base reconcilation trigger ignored, a reconcilation is already running")
    }
}

Cleaning up Knowledge Bases

func (m *KbManager) cleanupRemoteResources(ctx context.Context, args cleanupRemoteResourcesParams) error {
    err := m.awsClient.DeleteDataSource(ctx, &args.kbResourceId, &args.dataSourceId)
    if err != nil {
        return fmt.Errorf("error deleting data source remotely: %w", err)
    }

    _, err = m.awsClient.DeleteKnowledgeBase(ctx, args.kbResourceId)
    if err != nil {
        return fmt.Errorf("error deleting knowledge base remotely: %w", err)
    }

    err = m.awsClient.DeleteMultipleObjects(ctx, args.inclusionPrefix)
    if err != nil {
        return fmt.Errorf("error deleting knowledge base documents remotely: %w", err)
    }

    return nil
}

func (m *KbManager) CleanupKnowledgeBases(
    ctx context.Context,
) error {
    kbForCleanup, err := m.store.FetchKnowledgeBasesForCleanup(ctx)
    if err != nil {
        return fmt.Errorf("error fetching knowledge base cleaning up:%w", err)
    }

    if len(kbForCleanup) == 0 {
        m.baseLogger.Info().Msg("no knowledge bases to cleanup")
        return nil
    }

    m.baseLogger.Info().Msgf("found %d knowledge bases to cleanup", len(kbForCleanup))

    g, gCtx := errgroup.WithContext(ctx)

    g.SetLimit(m.maxConcurrentProvisioner)

    for _, kb := range kbForCleanup {

        kbToCleanup := kb

        g.Go(func() error {
            if err := gCtx.Err(); err != nil {
                m.baseLogger.Warn().Msgf("cleanup cancelled for knowledge base due to earlier error")
                return err
            }

            cleanupParams := cleanupRemoteResourcesParams{
                kbResourceId:    kbToCleanup.KnowledgeBaseResourceID.String,
                dataSourceId:    kbToCleanup.DataSourceID.String,
                inclusionPrefix: kbToCleanup.StoragePrefix,
            }
            err := m.cleanupRemoteResources(ctx, cleanupParams)
            if err != nil {
                m.baseLogger.Error().Err(err).Msg("failed to cleanup remote resources for knowledge bases")
            }

            _, err = m.store.DeleteAwsKnowledgeBase(ctx, kbToCleanup.ID)
            if err != nil {
                msg := "failed to update database status for cleanup"
                m.baseLogger.Error().Err(err).Msg(msg)
                return fmt.Errorf("%s: %w", msg, err)
            }

            return nil

        })
    }

    if err := g.Wait(); err != nil {
        m.baseLogger.Error().Err(err).Msg("cleanup cycle completed with one or more failures")
        return fmt.Errorf("concurrent cleanup failed: %w", err)
    }

    return nil
}

Similar to the index manager we cleanup remote resources, we mark Opensearch Index in Index Manager for cleanup and then delete the state in database for Knowledge base.
We also setup cron job for the same

func (m *KbManager) CleanupKnowledgeBaseWorker(ctx context.Context) {
    location, err := time.LoadLocation("America/Chicago")
    if err != nil {
        m.baseLogger.Error().Err(err).Msg("cannot load the timezone")
        return
    }

    c := cron.New(cron.WithLocation(location))

    schedule := "30 2 * * *"
    _, err = c.AddFunc(schedule, func() {
        m.baseLogger.Info().Str("schedule", schedule).Str("timezone", location.String()).Msg("triggered scheduled cleanup job")

        jobCtx, cancel := context.WithTimeout(context.Background(), 2*time.Hour)
        defer cancel()

        if err := m.CleanupKnowledgeBases(jobCtx); err != nil {
            m.baseLogger.Error().Err(err).Msg("scheduled cleanup job failed")
        } else {
            m.baseLogger.Info().Msg("scheduled cleanup job completed successfully")
        }
    })

    if err != nil {
        m.baseLogger.Error().Err(err).Msg("FATAL: could not schedule cleanup job")
        return
    }

    m.baseLogger.Info().Str("schedule", schedule).Str("timezone", location.String()).Msg("cleanup job scheduled successfully")
    c.Start()

    <-ctx.Done()

    m.baseLogger.Info().Msg("shutting down scheduled cleanup worker")

    stopCtx := c.Stop()
    <-stopCtx.Done()
    m.baseLogger.Info().Msg("cleanup scheduler stopped gracefully")

}

Running Knowledge base manager

func (m *KbManager) Run(ctx context.Context) error {
    m.baseLogger.Info().Msg("starting knowledge base manager workers")

    g, gCtx := errgroup.WithContext(ctx)

    g.Go(func() error {
        err := m.reconcilationWorker(gCtx)
        return err
    })

    g.Go(func() error {
        m.CleanupKnowledgeBaseWorker(gCtx)
        return gCtx.Err()
    })

    m.baseLogger.Info().Msg("all knowledge base manager workers are running")
    err := g.Wait()
    if err != nil && !errors.Is(err, context.Canceled) {
        m.baseLogger.Error().Err(err).Msg("knowledge base manager stopped due to an err")
        return err
    }

    m.baseLogger.Info().Msg("knowledge base manager stopped gracefully")
    return nil
}

As background goroutine we start reconciliation worker and cleanup worker.

Prime Pooling

Prime pooling is something very integral part of this orchestration drama because for the first time when application is running it won't have anything no pre-provision remote indexes or knowledge bases, so we wanted to initially pre-provision minimum pool requirement for those resources and then start the application.
We define method PrimePool for that.
We are having local context whose expiry is 60 minutes which is enough for this.
And we make sure all the indexes and knowledge bases are ready and not just warming up.

func (m *KbManager) PrimePool(ctx context.Context) error {
    m.baseLogger.Info().Msg("starting index pool priming...")

    indexReadyCtx, cancelIndex := context.WithTimeout(ctx, 60*time.Minute)
    defer cancelIndex()

    ticker := time.NewTicker(5 * time.Second)
    defer ticker.Stop()

    indexReady := false
    for !indexReady {
        select {
        case <-indexReadyCtx.Done():
            return fmt.Errorf("timeout waiting for index pool to be ready: %w", indexReadyCtx.Err())
        case <-ticker.C:
            hasEnoughIndexes, indexesWereProvisioned, err := m.indexManager.reconcileIndexPool(ctx)
            if err != nil {
                return fmt.Errorf("error reconciling index pool: %w", err)
            }
            if hasEnoughIndexes {
                if indexesWereProvisioned {
                    m.baseLogger.Info().Msg("new indexes were provisioned, waiting before reconciling knowledge bases")
                    time.Sleep(30 * time.Second)
                }
                m.baseLogger.Info().Msg("index pool has sufficient available indexes")
                indexReady = true
            }
        }
    }
    err := m.reconcileKnowledgeBasePool(ctx)
    if err != nil {
        return fmt.Errorf("error reconciling prime knowledge base pool: %w", err)
    }
    m.baseLogger.Info().Msg("knowledge base pool priming successful")
    return nil
}

Symphony - How it all come together to play out.

I have provided you main.go that's the entry point of application and this is how we handle things
We are creating context which gets canceled automatically when OS signals like pressing Ctrl + C (syscall.SIGINT), when systems sends terminating signal (syscall.SIGTERM).
Then we are creating logging object based on zerolog
we are loading our configuration from yaml file.
creating instance of database and AWS for their corresponding operation
Then we create group context for handling background goroutines
Then we initially synchronously make sure all the resources are present before starting the application using PrimePool

func main() {
    ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
    defer stop()

    logConfig := logger.LoggerConfig{
        LogLevel: zerolog.DebugLevel,
        FileConfig: &logger.FileConfig{
            Path:       "../logs/chatbot_engine/chatbot_engine.log",
            MaxSize:    10,
            MaxBackups: 5,
            MaxAge:     30,
            Compress:   true,
        },
    }

    baseLogger := logger.NewLogger(logConfig)

    config, err := utils.LoadConfig("./chatbot_engine.yaml")
    if err != nil {
        baseLogger.Fatal().Err(err).Msg("cannot load configuration")
    }

    switch config.Environment {
    case globals.ProdEnv:
        gin.SetMode(gin.ReleaseMode)
        baseLogger.Info().Msg("Gin mode set to 'release'")
    default:
        gin.SetMode(gin.DebugMode)
        baseLogger.Info().Msg("Gin mode set to 'debug'")
    }

    db, err := sql.Open("mysql", config.DBSource)
    if err != nil {
        baseLogger.Fatal().Err(err).Msg("error opening database connection")
    }
    defer db.Close()

    pingCtx, cancelPing := context.WithTimeout(ctx, 5*time.Second)
    defer cancelPing()
    err = db.PingContext(pingCtx)
    if err != nil {
        baseLogger.Fatal().Err(err).Msg("failed to connect to database (ping failed)")
    }

    db.SetMaxOpenConns(config.DBMaxOpenConn)
    db.SetMaxIdleConns(config.DBMaxIdleConn)
    db.SetConnMaxLifetime(time.Duration(config.DBConnMaxLifetime))

    baseLogger.Info().Msg("database connection pool initialized successfully")

    store := database.NewSqlStore(db)

    group, groupCtx := errgroup.WithContext(ctx)

    var awsClient *aws_cloud.AWSCloudClients
    if config.Environment != globals.DevEnv {
        var err error
        awsClient, err = aws_cloud.NewAWSCloudClients(
            groupCtx,
            aws_cloud.AwsCloudClientArg{
                DevEnvironmentCfg:  nil,
                Region:             config.AWSDefaultRegion,
                OpensearchEndpoint: config.AWSOpensearchEndpoint,
                BucketName:         config.AWSBucket,
                QueueName:          config.AWSProdQueueName,
                KeyID:              config.AWSKeyID,
            },
        )
        if err != nil {
            baseLogger.Fatal().Err(err).Msg("error creating aws client")
        }
    } else {
        var err error
        awsClient, err = aws_cloud.NewAWSCloudClients(
            groupCtx,
            aws_cloud.AwsCloudClientArg{
                DevEnvironmentCfg: &aws_cloud.DevEnvironmentConfig{
                    AccessKey: config.AWSAccessKey,
                    Secret:    config.AWSSecretAccess,
                },
                LogEnable:          config.AwsLogEnable,
                Region:             config.AWSDefaultRegion,
                OpensearchEndpoint: config.AWSOpensearchEndpoint,
                BucketName:         config.AWSBucket,
                QueueName:          config.AWSDevQueueName,
                KeyID:              config.AWSKeyID,
            },
        )
        if err != nil {
            baseLogger.Fatal().Err(err).Msg("error creating aws client")
        }
    }

    indexManager := knowledge_base.NewIndexManager(awsClient, store, baseLogger)
    kbManager := knowledge_base.NewKbManager(awsClient, store, baseLogger, &config, indexManager)

    if err := kbManager.PrimePool(groupCtx); err != nil {
        baseLogger.Fatal().Err(err).Msg("error prime pooling knowledge base and indexes")
    }

    group.Go(func() error {
        err := indexManager.Run(groupCtx)
        if err != nil {
            baseLogger.Error().Err(err).Msg("index manager exited with error")
        }

        return err
    })

    group.Go(func() error {
        err := kbManager.Run(groupCtx)
        if err != nil {
            baseLogger.Error().Err(err).Msg("knowledge base manager exited with error")
        }

        return err
    })

    err = group.Wait()
    if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, http.ErrServerClosed) {
        baseLogger.Error().Err(err).Msg("application exited due to error")
    } else {
        baseLogger.Info().Msg("application exiting gracefully")
    }

        baseLogger.Info().Msg("Bye :)")
}

This is how I worked around pre-provision, I hope this helps.

Thanks :)

Twitter: X

Building NeuroStash - VI

DEVunderdog — Sat, 09 Aug 2025 10:39:16 +0000

Raw thoughts from the trenches of building better retrieval systems

The Problem That Kept Me Up at Night

You know that feeling when your RAG system returns technically correct but completely useless chunks? Yeah, I was there. My traditional text splitters were butchering documents like a blunt knife through steak - creating arbitrary 512-token chunks that split sentences mid-thought, separated context from conclusions, and generally made my retrieval system feel like it had amnesia.

I was getting chunks like this:

Chunk 1: "...the economic impact was severe. Inflation rose to"
Chunk 2: "15% by the end of the quarter, causing widespread"
Chunk 3: "concern among investors and policymakers who..."

Useless. Absolutely useless.

Enter Semantic Chunking

While diving deep into LangChain's experimental modules, I stumbled upon something that made me question everything I thought I knew about text splitting. The SemanticChunker - a text splitter that actually understands what it's reading.

Instead of counting tokens like a robot, it uses embeddings to measure semantic similarity between sentences. Mind = blown.

How This Beast Actually Works

Let me break down the magic happening under the hood:

Step 1: Sentence Buffering (The Context Window Trick)

def combine_sentences(sentences: List[dict], buffer_size: int = 1) -> List[dict]:
    for i in range(len(sentences)):
        combined_sentence = ""

        # Add context BEFORE current sentence
        for j in range(i - buffer_size, i):
            if j >= 0:
                combined_sentence += sentences[j]["sentence"] + " "

        # Add current sentence
        combined_sentence += sentences[i]["sentence"]

        # Add context AFTER current sentence  
        for j in range(i + 1, i + 1 + buffer_size):
            if j < len(sentences):
                combined_sentence += " " + sentences[j]["sentence"]

        sentences[i]["combined_sentence"] = combined_sentence

    return sentences

This is brilliant. Instead of treating each sentence in isolation, it creates a sliding window of context. So instead of just embedding "The results were unexpected.", it embeds "The experiment ran for 6 months. The results were unexpected. Further analysis was needed."

Context matters. Always.

Step 2: Semantic Distance Calculation (Where the Magic Happens)

def calculate_cosine_distances(sentences: List[dict]) -> Tuple[List[float], List[dict]]:
    distances = []
    for i in range(len(sentences) - 1):
        embedding_current = sentences[i]["combined_sentence_embedding"]
        embedding_next = sentences[i + 1]["combined_sentence_embedding"]

        # Calculate semantic similarity
        similarity = cosine_similarity([embedding_current], [embedding_next])[0][0]

        # Convert to distance (higher = more different)
        distance = 1 - similarity
        distances.append(distance)
        sentences[i]["distance_to_next"] = distance

    return distances, sentences

Here's where it gets interesting. The system calculates how semantically different consecutive sentences are. High distance = topic shift. Low distance = same topic.

Step 3: Smart Breakpoint Detection

The system offers multiple strategies to decide where to split:

# Percentile approach - split at the top 5% most different transitions
breakpoint_threshold = np.percentile(distances, 95)

# Standard deviation - split when distance is 3 std devs above mean
breakpoint_threshold = np.mean(distances) + 3 * np.std(distances)

# Gradient approach - find where the rate of change spikes
distance_gradient = np.gradient(distances)
breakpoint_threshold = np.percentile(distance_gradient, 95)

I've been experimenting with all three, and honestly? Percentile at 95% has been my sweet spot for technical documentation.

Real-World Impact on My RAG Pipeline

Here's what changed when I swapped out my old chunker:

Before (Traditional Chunking):

# Old approach - dumb but fast
chunks = text_splitter.split_text(
    document, 
    chunk_size=512, 
    chunk_overlap=50
)

After (Semantic Chunking):

# New approach - smart and contextual
semantic_chunker = SemanticChunker(
    embeddings=OpenAIEmbeddings(),
    breakpoint_threshold_type="percentile",
    breakpoint_threshold_amount=95,
    buffer_size=1,
    min_chunk_size=100
)

chunks = semantic_chunker.split_text(document)

The Performance Trade-off (Because Nothing's Free)

Let's be real - this isn't free. Semantic chunking requires:

Embedding generation for every sentence group
Cosine similarity calculations
Statistical analysis for breakpoint detection

For a 10k word document:

Traditional chunking: ~50ms
Semantic chunking: ~2.3s

But here's the thing - those extra 2 seconds at indexing time save me hours of debugging why my RAG system returns garbage.

The Bottom Line

Semantic chunking isn't a silver bullet, but it's damn close to solving one of RAG's biggest problems. If you're building anything where context matters, you need to try this.

The code is all there in LangChain's experimental repo. Greg Kamradt deserves all the credit for the original implementation - I'm just the guy who got obsessed with making it work in production.

Building RAG systems that actually work, one semantic chunk at a time. Follow my journey at @DEVunderdog

Connect With Me: X

Building NeuroStash - V

DEVunderdog — Tue, 05 Aug 2025 18:49:15 +0000

Custom Exception Handling: From Chaos to Clean API Responses

In Part IV, I built a bulletproof document management system with state machines and atomic operations. But there was still one glaring problem: terrible error messages.

When users send malformed requests to your API, FastAPI's default error responses look like this:

{
  "detail": [
    {
      "type": "missing",
      "loc": ["body", "files"],
      "msg": "Field required",
      "input": {"wrong_field": "value"},
      "url": "https://errors.pydantic.dev/2.4/v/missing"
    },
    {
      "type": "string_too_short",
      "loc": ["body", "name"],
      "msg": "String should have at least 1 characters",
      "input": "",
      "url": "https://errors.pydantic.dev/2.4/v/string_too_short"
    }
  ]
}

This is developer-hostile. Frontend developers hate parsing nested arrays. Mobile apps crash trying to display this. Your API docs look unprofessional.

I needed something better. Something that makes sense.

The Problem: FastAPI's Default Validation Errors

FastAPI uses Pydantic for request validation, which is fantastic for type safety but terrible for API ergonomics. The default RequestValidationError responses are:

Inconsistent structure: Arrays instead of objects
Too verbose: Includes internal URLs and type information
Hard to parse: Nested location arrays like ["body", "files", 0, "name"]
Not user-friendly: Technical jargon instead of clear messages

When building NeuroStash's document upload system, I realized users would constantly hit validation errors:

Missing required fields in upload requests
Invalid file formats
Empty knowledge base names
Malformed JWT tokens

Each error needed to be clear, consistent, and actionable.

The Solution: Custom Exception Handler

I designed a custom exception handler that transforms FastAPI's validation chaos into clean, structured responses:

# app/core/exceptions.py
from fastapi import Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse

async def request_validation_exception_handler(
    request: Request, exc: RequestValidationError
):
    formatted_errors = {}

    for error in exc.errors():
        # Convert location array to dot notation
        field_name = ".".join(map(str, error["loc"]))
        message = error["msg"]
        formatted_errors[field_name] = message

    return JSONResponse(
        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
        content={
            "detail": {
                "message": "Request body validation failed. Please check the errors",
                "errors": formatted_errors,
            }
        },
    )

What This Handler Does

Flattens nested errors: Converts ["body", "files", 0] to "body.files.0"
Creates key-value pairs: Maps field names directly to error messages
Provides consistent structure: Always returns the same response format
Maintains helpful context: Includes both summary message and specific errors

Before vs After: The Transformation

Before (FastAPI default):

{
  "detail": [
    {
      "type": "missing",
      "loc": ["body", "files"],
      "msg": "Field required",
      "input": {"name": "My KB"},
      "url": "https://errors.pydantic.dev/2.4/v/missing"
    },
    {
      "type": "string_too_short", 
      "loc": ["body", "name"],
      "msg": "String should have at least 1 characters",
      "input": "",
      "url": "https://errors.pydantic.dev/2.4/v/string_too_short"
    }
  ]
}

After (Custom handler):

{
  "detail": {
    "message": "Request body validation failed. Please check the errors",
    "errors": {
      "body.files": "Field required",
      "body.name": "String should have at least 1 characters"
    }
  }
}

The transformation is dramatic. Frontend developers can now:

Access errors directly: response.detail.errors["body.files"]
Display user-friendly messages without parsing
Build consistent error handling across all endpoints

Integration with NeuroStash

I registered the custom handler in the main application:

# main.py
from fastapi import FastAPI
from fastapi.exceptions import RequestValidationError
from app.core.exceptions import request_validation_exception_handler

app = FastAPI(title=settings.PROJECT_NAME)

# Register the custom exception handler
app.add_exception_handler(
    RequestValidationError, 
    request_validation_exception_handler
)

This single line transforms every validation error across the entire application.

Real-World Impact

Here's how this improved NeuroStash's API experience:

Document Upload Errors

When users forget required fields in document uploads:

# POST /documents/upload
{
  "wrong_field": ["file1.pdf", "file2.docx"]  # Missing 'files' field
}

Response:

{
  "detail": {
    "message": "Request body validation failed. Please check the errors",
    "errors": {
      "body.files": "Field required"
    }
  }
}

Knowledge Base Creation Errors

When knowledge base names are empty:

# POST /kb/create  
{
  "name": ""  # Too short
}

Response:

{
  "detail": {
    "message": "Request body validation failed. Please check the errors", 
    "errors": {
      "body.name": "String should have at least 1 characters"
    }
  }
}

Complex Nested Validation

For complex objects with nested validation errors:

# POST /documents/finalize
{
  "successful": "not-an-array",  # Wrong type
  "failed": [1, 2, "not-a-number"]  # Mixed types
}

Response:

{
  "detail": {
    "message": "Request body validation failed. Please check the errors",
    "errors": {
      "body.successful": "Input should be a valid list",
      "body.failed.2": "Input should be a valid integer"
    }
  }
}

Advanced Pattern: Field-Specific Improvements

The dot notation makes it easy to handle complex validation scenarios:

def get_user_friendly_message(field_name: str, original_message: str) -> str:
    """Convert technical validation messages to user-friendly ones"""

    field_mappings = {
        "body.files": "Please provide a list of files to upload",
        "body.name": "Knowledge base name cannot be empty", 
        "body.successful": "Successful uploads must be a list of document IDs",
        "body.failed": "Failed uploads must be a list of document IDs"
    }

    return field_mappings.get(field_name, original_message)

Error Handling Strategy

This custom handler is part of a broader error handling strategy in NeuroStash:

Validation Errors (422): Handled by custom exception handler
Business Logic Errors (409, 404): Custom exceptions with specific messages
Authentication Errors (401, 403): JWT middleware with clear responses
Server Errors (500): Logged with correlation IDs, generic user message

Each type of error gets appropriate treatment while maintaining consistent response structure.

Best Practices Learned

1. Consistent Response Structure

Every error response follows the same pattern:

{
  "detail": {
    "message": "Human-readable summary",
    "errors": { /* field-specific errors */ }
  }
}

2. Frontend-Friendly Field Names

Dot notation (body.files.0.name) is easy to parse and display in UI forms.

3. Preserve HTTP Status Codes

Keep FastAPI's default status codes (422 for validation) so HTTP clients behave correctly.

4. Don't Leak Internal Details

Remove Pydantic URLs and type information that users don't need.

5. Make Errors Actionable

Each error message should tell users what to fix, not just what's wrong.

Key Takeaways

Default FastAPI validation errors are developer-hostile - fix them early
Consistent error structure makes frontend integration painless
Dot notation field names are easy to parse and display
Custom exception handlers transform your entire API with one registration
Think about your API consumers - clean errors improve developer experience

Building production APIs isn't just about functionality - it's about creating interfaces that developers actually want to use. Clean error handling is a cornerstone of good API design.

Want to see the full implementation? Check out the NeuroStash repository and follow along as we build something amazing! 🚀

Connect with me: @fkhoja098

Building NeuroStash - IV

DEVunderdog — Thu, 31 Jul 2025 19:25:53 +0000

Part 4 of the NeuroStash series - Designing a bulletproof document management system with knowledge base integration

In Part III, I covered the RBAC system that controls who can do what in NeuroStash. Now comes the exciting part - building the core functionality that users actually interact with: document management and knowledge base creation.

This isn't just about storing files. Document management requires handling:

Atomic uploads with proper rollback mechanisms
Conflict resolution when uploads fail midway
Presigned URL generation for secure direct-to-S3 uploads
Document lifecycle management with proper state tracking
Knowledge base associations with many-to-many relationships

Let me walk you through how I solved these challenges.

The Document Upload Challenge

Most engineers think document upload is simple: "Just accept the file and save it." But in production, you quickly discover the edge cases:

What if the upload fails halfway through?
What if the same filename is uploaded twice?
How do you prevent orphaned database records?
How do you handle large files without timeouts?

I designed a two-phase upload system that solves all these problems.

Phase 1: Presigned URL Generation

Instead of uploading through the API server (which creates bottlenecks), I generate presigned URLs for direct S3 uploads:

@router.post("/upload", response_model=GeneratedPresignedUrls)
def upload_documents(
    req: GeneratePresignedUrlsReq,
    db: SessionDep,
    payload: TokenPayloadDep,
    aws_client: AwsDep,
):
    list_of_documents: List[CreateDocument] = []
    url_by_filename: Dict[str, str] = {}

    for file in req.files:
        # Generate unique filename to prevent conflicts
        unique_id = uuid.uuid4()
        filename = f"{file}-{unique_id}"
        object_key = f"{payload.user_id}/{filename}"

        # Generate presigned URL for direct S3 upload
        content_type = aws_client.extract_content_type(filename=file)
        url = aws_client.generate_presigned_upload_url(
            object_key=object_key, content_type=content_type
        )

        url_by_filename[filename] = url
        document = CreateDocument(
            user_id=payload.user_id, 
            file_name=filename, 
            object_key=object_key
        )
        list_of_documents.append(document)

    # Create database records in PENDING state
    created_documents = create_document(db=db, files=list_of_documents)

    final_response: Dict[int, str] = {}
    for doc_id, filename in created_documents:
        presigned_url = url_by_filename.get(filename)
        if presigned_url:
            final_response[doc_id] = presigned_url

    return GeneratedPresignedUrls(
        message="generated presigned urls successfully", 
        urls=final_response
    )

Key Design Decisions:

Unique filenames: Append UUID to prevent conflicts
User-scoped object keys: Isolate user data with {user_id}/{filename}
PENDING state: Documents start locked until upload completes
Atomic database operations: All records created in single transaction

The Database Schema: Built for Reliability

The document registry schema handles the complex state management:

class DocumentRegistry(Base, TimestampMixin):
    __tablename__ = "documents_registry"

    id: Mapped[int] = mapped_column(Integer, Identity(), primary_key=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("user_clients.id", onupdate="CASCADE", ondelete="RESTRICT"),
        nullable=False,
    )
    file_name: Mapped[str] = mapped_column(String(100), nullable=False)
    object_key: Mapped[str] = mapped_column(String(150), nullable=False)
    lock_status: Mapped[bool] = mapped_column(Boolean, nullable=False)
    op_status: Mapped[OperationStatusEnum] = mapped_column(
        SQLEnum(OperationStatusEnum),
        nullable=False,
        server_default=OperationStatusEnum.PENDING.value,
    )

    __table_args__ = (
        UniqueConstraint("user_id", "file_name", name="idx_unique_filename"),
        Index("idx_file_registry_user_id", "user_id", "lock_status", "op_status"),
    )

The idea of the state machine:

PENDING + locked: Upload in progress
SUCCESS + unlocked: Ready for use
FAILED + unlocked: Upload failed, can retry
PENDING + unlocked: Conflict state (needs cleanup)

Phase 2: Upload Finalization

After S3 upload completes (or fails), the client calls the finalize endpoint:

@router.put("/finalize", response_model=StandardResponse)
def post_upload_documents(req: FinalizeDocumentReq, db: SessionDep):
    try:
        finalize_documents(
            db=db, 
            successful=req.successful, 
            failed=req.failed
        )
        return StandardResponse(message="successfully finalized documents")
    except Exception:
        logger.exception("error finalizing document")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="error finalizing documents",
        )

what is actually the finalize_documents function about ? So here it is,

def finalize_documents(*, db: Session, successful: List[int], failed: List[int]):
    try:
        all_ids = successful + failed

        stmt = (
            update(DocumentRegistry)
            .where(DocumentRegistry.id.in_(all_ids))
            .values(
                op_status=case(
                    (
                        DocumentRegistry.id.in_(successful),
                        cast(
                            OperationStatusEnum.SUCCESS.value,
                            DocumentRegistry.op_status.type,
                        ),
                    ),
                    (
                        DocumentRegistry.id.in_(failed),
                        cast(
                            OperationStatusEnum.FAILED.value,
                            DocumentRegistry.op_status.type,
                        ),
                    ),
                ),
                lock_status=False,
            )
        )

        db.execute(stmt)
        db.commit()
    except Exception:
        db.rollback()
        raise

This single query updates hundreds of documents atomically - successful ones get SUCCESS status, failed ones get FAILED status, and all get unlocked.

Knowledge Base Creation: Simple but Robust

@router.post("/create", response_model=CreatedKb)
def create_knowledge_base(
    req: CreateKbReq, 
    db: SessionDep, 
    token_payload: TokenPayloadDep
):
    try:
        args = CreateKbInDb(user_id=token_payload.user_id, name=req.name)
        created_kb = create_kb_db(db=db, kb=args)
        return CreatedKb(
            message="successfully created knowledge base",
            id=created_kb.id,
            name=created_kb.name,
        )
    except KnowledgeBaseAlreadyExists:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT, 
            detail="knowledge base already exists"
        )

The database function handles the uniqueness constraint:

def create_kb_db(*, db: Session, kb: CreateKbInDb) -> KnowledgeBase:
    try:
        knowledge_base = KnowledgeBase(**kb.model_dump())
        db.add(knowledge_base)
        db.commit()
        db.refresh(knowledge_base)
        return knowledge_base
    except IntegrityError as e:
        db.rollback()
        if isinstance(e.orig, psycopg.errors.UniqueViolation):
            raise KnowledgeBaseAlreadyExists(knowledg_base_name=kb.name)
        else:
            raise RuntimeError(f"database integrity error: {e}")

Custom exceptions make error handling clean - the API layer can catch specific exceptions and return appropriate HTTP status codes.

Knowledge Base Architecture: Many-to-Many Done Right

Documents alone aren't useful - you need to organize them into knowledge bases. This requires a many-to-many relationship where:

One document can belong to multiple knowledge bases
One knowledge base can contain multiple documents
Each association has its own status tracking

The Junction Table Pattern

class KnowledgeBaseDocument(Base, TimestampMixin):
    __tablename__ = "knowledge_base_documents"

    id: Mapped[int] = mapped_column(Integer, Identity(), primary_key=True)
    knowledge_base_id: Mapped[int] = mapped_column(
        ForeignKey("knowledge_bases.id", onupdate="CASCADE", ondelete="CASCADE"),
        nullable=False,
    )
    document_id: Mapped[int] = mapped_column(
        ForeignKey("documents_registry.id", onupdate="CASCADE", ondelete="CASCADE"),
        nullable=False,
    )
    status: Mapped[OperationStatusEnum] = mapped_column(
        SQLEnum(OperationStatusEnum),
        nullable=False,
        server_default=OperationStatusEnum.PENDING.value,
    )

    __table_args__ = (
        UniqueConstraint(
            "knowledge_base_id", "document_id", 
            name="idx_unique_kb_doc_combination"
        ),
    )

Why not just use a simple junction table? Because, we are going to have ingestion job which is going to be asynchronous and having a state machine helps to track things and clean them up. A document might be:

PENDING: Being processed for the knowledge base
SUCCESS: Successfully added and indexed
FAILED: Processing failed, needs retry

Document Lifecycle: The Cleanup Problem

In production, things go wrong. Networks fail, users close browsers, S3 timeouts occur. This creates "conflicted" documents - database records that don't match S3 reality.

I built an automated cleanup system that is ofcourse triggered by API endpoint:

@router.get("/cleanup", response_model=StandardResponse)
def cleanup_files(db: SessionDep, payload: TokenPayloadDep, aws_client: AwsDep):
    conflicting_docs = conflicted_docs(db=db, user_id=payload.user_id)

    if not conflicting_docs:
        return StandardResponse(message="no conflicting files found")

    to_be_unlocked = []
    to_be_deleted = []

    for doc in conflicting_docs:
        exists: bool = aws_client.object_exists(object_key=doc.object_key)
        if not exists:
            # Database record but no S3 object - delete record
            to_be_deleted.append(doc.id)
        else:
            # S3 object exists but record is conflicted - fix status
            to_be_unlocked.append(doc.id)

    if to_be_deleted or to_be_unlocked:
        cleanup_docs(
            db=db,
            user_id=payload.user_id,
            to_be_unlocked=to_be_unlocked,
            to_be_deleted=to_be_deleted,
        )

    return StandardResponse(message="cleanup successfully finished")

The conflicted_docs function finds documents in conflicted states:

def conflicted_docs(*, db: Session, user_id: int) -> List[DocumentRegistry]:
    # Valid state combinations
    valid_combinations = [
        (True, OperationStatusEnum.PENDING),   # Upload in progress
        (True, OperationStatusEnum.SUCCESS),   # Processing
        (True, OperationStatusEnum.FAILED),    # Failed processing
        (False, OperationStatusEnum.PENDING),  # Failed upload
        (False, OperationStatusEnum.FAILED),   # Available failed
        # (False, OperationStatusEnum.SUCCESS) - This is the only valid "ready" state
    ]

    stmt = select(DocumentRegistry).where(
        and_(
            DocumentRegistry.user_id == user_id,
            or_(
                *[
                    and_(
                        DocumentRegistry.lock_status == lock_status,
                        DocumentRegistry.op_status == op_status,
                    )
                    for lock_status, op_status in valid_combinations
                ]
            ),
        )
    )

    result = db.execute(stmt)
    return result.scalars().all()

This finds documents that need attention and reconciles database state with S3 reality.

Listing and Pagination: Performance Matters

When you have thousands of documents, naive listing kills performance. I implemented efficient pagination:

def list_files(
    *, db: Session, user_id: int, limit: int, offset: int
) -> Tuple[List[DocumentRegistry], int]:
    try:
        # Base query for ready documents only
        stmt = select(DocumentRegistry).where(
            DocumentRegistry.user_id == user_id,
            DocumentRegistry.lock_status == False,
            DocumentRegistry.op_status == OperationStatusEnum.SUCCESS,
        )

        # Count total for pagination metadata
        count_stmt = select(func.count()).select_from(stmt.subquery())
        total_count = db.execute(count_stmt).scalar()

        # Apply pagination
        stmt = stmt.limit(limit=limit).offset(offset=offset)
        documents = db.execute(stmt).scalars().all()

        return documents, total_count
    except Exception:
        logger.error("error listing user documents")
        raise

Key optimizations:

Filter at database level: Only query documents in SUCCESS state
Separate count query: Get total count for pagination metadata
Index support: The composite index on (user_id, lock_status, op_status) makes this blazing fast

Real-World Usage Flow

Here's how the complete system works in practice:

# 1. Request presigned URLs for upload
POST /documents/upload
Authorization: Bearer <jwt-token>
{
  "files": ["document1.pdf", "document2.docx"]
}

# Returns:
{
  "urls": {
    "123": "https://s3.amazonaws.com/bucket/presigned-url-1",
    "124": "https://s3.amazonaws.com/bucket/presigned-url-2"
  }
}

# 2. Upload directly to S3 using presigned URLs
PUT https://s3.amazonaws.com/bucket/presigned-url-1
Content-Type: application/pdf
<binary data>

# 3. Finalize the upload status
PUT /documents/finalize
{
  "successful": [123],
  "failed": [124]
}

# 4. Create knowledge base
POST /kb/create
{
  "name": "Engineering Docs"
}

# 5. List available documents
GET /documents/list?limit=20&offset=0

# 6. Add documents to knowledge base
# (This would be in the ingestion system - Part V!)

Lessons Learned & Best Practices

State Machines Save Your Sanity
Don't use boolean flags for complex states. Enums with clear state transitions make debugging infinitely easier.
Always Plan for Failure
Networks fail, uploads timeout, users close browsers. Design your system to handle partial failures gracefully.
Presigned URLs Scale Better
Direct-to-S3 uploads bypass your API servers entirely, eliminating bottlenecks and reducing costs.
Cleanup is Not Optional
In production, orphaned records accumulate. Build cleanup mechanisms from day one.
Index Your Queries
The composite index on (user_id, lock_status, op_status) makes document listing sub-millisecond even with millions of records.

Want to see the full implementation? Check out the NeuroStash repository and follow along as we build something amazing! 🚀

Connect with me: X

Building NeuroStash - III

DEVunderdog — Tue, 29 Jul 2025 18:58:00 +0000

Part 3 of the NeuroStash series - Building an enterprise-grade knowledge management platform

In my previous post, I walked through building a robust authentication system for NeuroStash using JWT tokens, API keys, and AWS KMS encryption. Today, I'm diving deeper into what happens after authentication - how I designed role-based access control (RBAC) and user management.

The User Management Challenge

After solving authentication, I faced a new set of questions:

How do you onboard new users securely?
What happens when someone needs elevated permissions?
How do you prevent privilege escalation attacks?
How do you make role management intuitive for administrators?

Most engineers fall into two traps:

Over-engineering: Complex permission matrices that become unmaintainable
Under-engineering: Simple boolean flags that don't scale

I needed something in between - powerful enough for enterprise needs, simple enough to reason about.

Architecture Overview: The Foundation

Building on the authentication system from Part II, here's how the RBAC system works:

# The core role enum - keeping it simple but extensible
class ClientRoleEnum(str, Enum):
    USER = "user"
    ADMIN = "admin"

Why only two roles? In my experience, most SaaS products start with user/admin and only add complexity when needed. NeuroStash follows this principle - you can always extend later, but you can't easily simplify.

User Registration: Admin-Controlled Onboarding

Instead of public registration (a security nightmare for enterprise), I implemented admin-controlled user onboarding:

@router.post("/register", response_model=UserClientCreated)
def register_user_to_app(
    user_in: RegisterUser,
    db: SessionDep,
    token_manager: TokenDep,
    admin_payload: TokenPayloadDep,  # Admin verification
):
    # First, verify the requester is an admin
    if admin_payload.role != ClientRoleEnum.ADMIN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="you are not authorized to perform this action",
        )

    # Generate API key for the new user
    api_key, api_key_bytes, api_key_signature, active_key_id = (
        token_manager.generate_api_key()
    )

    # Create user with default USER role
    user = UserClientCreate(email=user_in.email, role=ClientRoleEnum.USER)
    api_key_params = ApiKeyCreate(
        key_id=active_key_id,
        key_credential=api_key_bytes,
        key_signature=api_key_signature,
    )

    db_user_client, db_api_key = register_user(
        db=db, user=user, api_key_params=api_key_params
    )

    return UserClientCreated(
        id=db_user_client.id, 
        email=db_user_client.email, 
        api_key=api_key  # Returned to admin for sharing
    )

Key Design Decisions:

Admin-only registration: Prevents unauthorized signups
Automatic API key generation: Every user gets immediate API access
Default USER role: Principle of least privilege
API key returned to admin: Secure key distribution workflow

Role Promotion: Controlled Privilege Escalation

Promoting users to admin requires careful handling:

@router.patch("/promote/{user_id}", response_model=StandardResponse)
def promote_users(
    user_id: int,
    db: SessionDep,
    admin_payload: TokenPayloadDep,
):
    # Verify admin permissions
    if admin_payload.role != ClientRoleEnum.ADMIN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="you are not authorized to perform this action",
        )

    # Validate user_id
    if user_id == 0:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="please provide user_id to promote",
        )

    # Perform the promotion
    user_client = promote_user_db(db=db, user_id=user_id)
    if user_client is None:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="cannot find user with provided id",
        )

    return StandardResponse(message="successfully promoted user to admin")

The database logic handles idempotency elegantly:

def promote_user_db(*, db: Session, user_id: int) -> UserClient | None:
    stmt = select(UserClient).where(UserClient.id == user_id)
    user_client = db.scalars(stmt).first()

    if user_client:
        # Check if already admin (idempotent operation)
        if user_client.role == ClientRoleEnum.ADMIN:
            return user_client

        # Promote to admin
        user_client.role = ClientRoleEnum.ADMIN
        try:
            db.commit()
            return user_client
        except Exception:
            db.rollback()
            raise
    else:
        return None

User Deletion: Preventing Self-Destruction

Deleting users requires protection against admin self-deletion:

@router.delete("/delete/{user_id}", response_model=StandardResponse)
def delete_users(
    user_id: int,
    db: SessionDep,
    admin_payload: TokenPayloadDep,
):
    # Standard admin check
    if admin_payload.role != ClientRoleEnum.ADMIN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="you are not authorized to perform this action",
        )

    # Prevent self-deletion (critical!)
    if admin_payload.user_id == user_id:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, 
            detail="you cannot delete yourself"
        )

    # Proceed with deletion
    deleted = delete_user_db(db=db, user_id=user_id)
    if not deleted:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="cannot find user with provided id",
        )

    return StandardResponse(message="user deleted successfully")

The self-deletion check is crucial - I've seen production systems where admins accidentally locked themselves out. This simple check prevents that nightmare scenario.

Dependency Injection: Clean Security Boundaries

FastAPI's dependency injection makes role checks clean and reusable:

# Clean type annotations for dependencies
TokenPayloadDep = Annotated[TokenData, Depends(get_token_payload)]
ApiPayloadDep = Annotated[ApiData, Depends(get_api_payload)]

# Usage in routes
@router.get("/admin-only-endpoint")
def admin_only_route(payload: TokenPayloadDep):
    if payload.role != ClientRoleEnum.ADMIN:
        raise HTTPException(status_code=401, detail="admin required")

    # Admin logic here
    return {"message": "admin access granted"}

This pattern offers several benefits:

Consistent authentication: Every route gets the same validation
Type safety: PayloadDep ensures you have user context
Separation of concerns: Auth logic stays in dependencies

Database Schema: The Foundation Layer

The user schema supports the RBAC system with careful constraints:

class UserClient(Base, TimestampMixin):
    __tablename__ = "user_clients"

    id: Mapped[int] = mapped_column(Integer, Identity(), primary_key=True)
    email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
    role: Mapped[ClientRoleEnum] = mapped_column(
        Enum(ClientRoleEnum), 
        nullable=False, 
        default=ClientRoleEnum.USER
    )

    # Relationship to API keys
    api_keys: Mapped[List["ApiKey"]] = relationship(
        "ApiKey", 
        back_populates="user_client",
        cascade="all, delete-orphan" 
    )

Key design elements:

Unique email constraint: Prevents duplicate accounts
Default USER role: Secure by default
Cascade deletion: API keys are cleaned up automatically
Timestamp mixin: Audit trail for user operations

User Listing: Admin Visibility with Pagination

Admins need visibility into the user base:

@router.get("/list", response_model=ListUsers)
def list_users(
    admin_payload: TokenPayloadDep,
    db: SessionDep,
    limit: int = 10,
    offset: int = 0,
):
    if admin_payload.role != ClientRoleEnum.ADMIN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="you are not authorized to perform this action",
        )

    users = list_users_db(db=db, limit=limit, offset=offset)
    return ListUsers(message="successfully fetched users", users=users)

The database function handles pagination efficiently:

def list_users_db(*, db: Session, limit: int = 10, offset: int = 0) -> List[UserClient]:
    stmt = select(UserClient).order_by(UserClient.id).limit(limit).offset(offset)
    users = db.scalars(stmt).all()
    return users

Why pagination matters: Even with hundreds of users, the API stays responsive. The order_by(UserClient.id) ensures consistent ordering across pages.

Integration with Authentication System

The RBAC system builds seamlessly on the authentication foundation from Part II:

# From authentication (Part II)
async def get_token_payload(
    token: Annotated[Optional[str], Depends(oauth2_scheme)],
    token_manager: TokenDep,
) -> TokenData:
    # JWT verification logic...
    return TokenData(email=payload.email, user_id=payload.user_id, role=payload.role)

This layered approach separates authentication (who are you?) from authorization (what can you do?).

Real-World Usage Examples

Here's how the complete system works in practice:

# 1. Admin registers a new user
POST /user/register
Authorization: Bearer <admin-jwt-token>
{
  "email": "newuser@company.com"
}
# Returns: API key for the new user

# 2. New user generates their own JWT token
GET /auth/generate/token
Authorization: ApiKey <user-api-key>
# Returns: JWT token for session use

# 3. Admin promotes user to admin
PATCH /user/promote/123
Authorization: Bearer <admin-jwt-token>
# User 123 now has admin privileges

# 4. Admin lists all users
GET /user/list?limit=20&offset=0
Authorization: Bearer <admin-jwt-token>
# Returns: Paginated user list

Lessons Learned & Best Practices

Start Simple, Scale Thoughtfully Two roles (USER/ADMIN) cover 90% of use cases. Don't build complex permission systems until you need them.
Admin Self-Destruction Prevention Always check if an admin is trying to delete/demote themselves. This simple check prevents lockout scenarios.
Idempotent Operations Promoting an admin should succeed, not fail. Design operations to be safe to retry.
Consistent Error Messages Standardize authorization error messages. Don't leak internal system details.
Dependency Injection for Security FastAPI's dependency system makes security checks clean

Want to see the full implementation? Check out the NeuroStash repository: Code

Building Neurostash - II

DEVunderdog — Mon, 28 Jul 2025 10:47:42 +0000

How I build Authentication for NeuroStash

When I started building NeuroStash, my AI-powered knowledge management platform, I knew authentication would be critical. Not just "password and done" authentication, but a better security that could handle API keys, JWT tokens.

Here’s how I architected a robust authentication system using Python, FastAPI, AWS KMS and PostgreSQL

The Authentication Challenge

Most developers face a common dilemma with authentication:

Too simple: Basic username/password systems are vulnerable and don't scale
Too complex: Over-engineered solutions become maintenance nightmares

I needed something that was both secure and maintainable, supporting:

JWT tokens for session-based authentication
API keys for programmatic access
Symmetric Encryption Keys that signs the JWT Token and API Keys stored in Database
Those Symmetric Encryption Keys which signs are needed to be encrypted via AWS KMS that produces double layer of security in case if database breach you can’t reverse engineer the keys.
We have Role-Base access control

Architecture Overview

AWS integration with smart resource management

The foundation starts with AWS KMS integration. I used Python's @property decorator to implement lazy loading of AWS clients

class AwsClientManager:
    @property
    def kms(self):
        if self._kms_client is None:
            self._kms_client = self.session.client("kms")
        return self._kms_client

Why this pattern works:

KMS clients are only created when needed
Reduces startup time and memory usage
Makes testing easier (you can mock the property)
Handles AWS credential rotation gracefully

The encrypt/decrypt methods wrap KMS operations with proper error handling

def encrypt_key(self, key_blob: bytes) -> Optional[bytes]:
        if not self.kms or not self.kms_key_id:
            logger.error("kms client or kms key id not configured for encryption")
            return None
        try:
            response = self.kms.encrypt(
                KeyId=self.settings.AWS_KMS_KEY_ID, Plaintext=key_blob
            )
            return response.get("CiphertextBlob")
        except ClientError as e:
            logger.error(f"error encrypting key with kms key id: {e}")
            raise

    def decrypt_key(self, ciphertext_blob: bytes) -> Optional[bytes]:
        if not self.kms or not self.kms_key_id:
            logger.error("kms client or kms key id not configured for encryption")
            return None
        try:
            response = self.kms.decrypt(
                CiphertextBlob=ciphertext_blob, KeyId=self.settings.AWS_KMS_KEY_ID
            )
            return response.get("Plaintext")
        except ClientError as e:
            logger.error(f"error decrypting key: {e}")
            raise

Symmetric Key Generation

For cryptographic operations, I used Python’s secrets module

def generate_symmetric_key() -> bytes:
    try:
        key = secrets.token_bytes(32)  # 256-bit key
        return key
    except Exception as e:
        logger.error(f"failed to generate key: {e}")
        raise KeyGenerationError(f"failed to generate key: {e}") from e

Why 256-bit keys: They provide excellent security while maintaining good performance for HMAC operations.

Database Schema Design

class EncryptionKey(Base, TimestampMixin):
    __tablename__ = "encryption_keys"

    id: Mapped[int] = mapped_column(Integer, Identity(), primary_key=True)
    symmetric_key: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
    is_active: Mapped[bool] = mapped_column(Boolean, nullable=False)
    expired_at: Mapped[Optional[datetime]] = mapped_column(TIMESTAMP(timezone=True))

    __table_args__ = (
        Index("idx_encryption_keys_active", "id", 
              postgresql_where=Column("is_active")),
    )

Performance optimization: The partial index on is_active means queries for active keys are lightning fast, even with thousands of expired keys in the table.

Token Manager: The heart of authentication

The TokenManager class orchestrates all authentication operations. It maintains an in-memory cache of decrypted keys

def _build_active_key_tuple(self, db: Session) -> Tuple[Dict[int, KeyInfo], int]:
    # Fetch encrypted keys from database
    active_encryption_keys = get_active_encryption_key(db=db)
    other_encryption_keys = get_other_encryption_keys(db=db)

    # Decrypt and cache in memory
    for key_id, value in key_info.items():
        decrypted_key_bytes = self._aws_client_manager.decrypt_key(value.key)
        decrypted_key_info[key_id] = KeyInfo(
            key=decrypted_key_bytes, 
            expires_at=value.expires_at
        )

    return (decrypted_key_info, active_id)

This approach gives us:

Fast token verification: No KMS calls during verification
Graceful key rotation: Multiple keys can be active simultaneously
Proper error handling: Failed decryption is caught early

JWT Token Implementation

JWT tokens include kid (key id) header for key identification and then that key which gets identified is used to sign it.

def create_access_token(self, payload_data: TokenData) -> str:
    all_keys, active_key_id = self.get_keys()
    active_key_info = all_keys[active_key_id]

    to_encode = payload_data.model_dump(mode='json', exclude_unset=True)
    to_encode.update({
        "exp": expire,
        "iss": settings.JWT_ISSUER,
        "aud": settings.JWT_AUDIENCE,
        "iat": datetime.now(timezone.utc),
        "jti": os.urandom(16).hex(),  # Unique token ID
    })

    headers = {"kid": active_key_id}
    return jwt.encode(to_encode, active_key_info.key.hex(), 
                     algorithm=ALGORITHMS.HS256, headers=headers)

Verification: extracts the kid and uses corresponding keys

def verify_token(self, token: str) -> Optional[TokenData]:
    unverified_headers = jwt.get_unverified_headers(token=token)
    kid = unverified_headers.get("kid")

    all_keys, _ = self.get_keys()
    key_for_verification = all_keys.get(kid)

    if key_for_verification.is_expired():
        raise RuntimeError(f"key {kid} has expired")

    payload = jwt.decode(token=token, key=key_for_verification.key.hex(),
                        algorithms=[ALGORITHMS.HS256])
    return TokenData(**payload)

API Key System: Stateless and Secure

API keys use HMAC signature for tamper-proof verification

def generate_api_key(self) -> Tuple[str, bytes, bytes, int]:
    # Generate random component
    random_bytes = secrets.token_bytes(24)
    random_bytes_b64 = base64.urlsafe_b64encode(random_bytes).decode("utf-8").rstrip("=")

    # Create HMAC signature
    data_to_hmac = f"{active_key_id}:{random_bytes_b64}".encode("utf-8")
    hmac_obj = hmac.new(active_key_info.key, data_to_hmac, hashlib.sha256)
    signature_bytes = hmac_obj.digest()
    signature_b64 = base64.urlsafe_b64encode(signature_bytes).decode("utf-8").rstrip("=")

    api_key = f"{random_bytes_b64}.{signature_b64}"
    return api_key, api_key_bytes, signature_bytes, active_key_id

def verify_api_key(self, api_key: str, key_hmac: bytes, kid: int) -> bool:
    parts = api_key.split(".")
    random_bytes_b64, signature_b64 = parts

    # Recreate HMAC with our key
    data_to_hmac = f"{kid}:{random_bytes_b64}".encode("utf-8")
    expected_hmac_obj = hmac.new(key_info.key, data_to_hmac, hashlib.sha256)
    expected_signature_bytes = expected_hmac_obj.digest()

    # Use constant-time comparison
    return hmac.compare_digest(expected_signature_bytes, client_signature_bytes)

Security Benefits:

Tamper-proof with HMAC-SHA256
Constant-time comparison prevents timing attacks
URL-safe Base64 encoding

FastAPI Integration

FastAPI's dependency injection makes authentication clean

async def get_token_payload(
    token: Annotated[Optional[str], Depends(oauth2_scheme)],
    token_manager: TokenDep,
) -> TokenData:
    if token is None:
        raise HTTPException(status_code=401, detail="missing bearer token")

    try:
        payload = token_manager.verify_token(token=token.credentials)
        return payload
    except ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="token has expired")
    except KeyNotFoundError as e:
        raise HTTPException(status_code=401, detail=f"key error: {e}")

Routes can simply depend on TokenPayloadDep or ApiPayloadDep

@router.post("/protected-endpoint")
def protected_route(payload: TokenPayloadDep):
    # payload.user_id, payload.email, payload.role are available
    return {"message": f"Hello {payload.email}"}

Lessons Learned

The @property decorator is underrated: Perfect for lazy resource initialization
HMAC-based API keys scale better: Stateless verification beats database lookups
FastAPI dependencies simplify auth: Clean separation of concerns
Error handling matters: Fail gracefully without information leakage

Wrapping Up

Building authentication is challenging, but breaking it into focused components makes it manageable. The combination of AWS KMS, careful database design, and FastAPI's dependency system creates a robust foundation.

The key is balancing security with performance and maintainability. This architecture handles enterprise requirements while keeping the codebase clean and testable.

Want to see the full implementation? Check out the NeuroStash repository: https://github.com/DEVunderdog/NeuroStash

Building Neurostash - I

DEVunderdog — Sun, 27 Jul 2025 05:49:45 +0000

How I’m architecting and building a scalable, multi-tenant document processing system.

The Problem That Started It All

We've all been there—drowning in documents scattered across emails, cloud drives, and local folders. Important information buried in PDFs, critical insights lost in meeting notes, and that one crucial document that you know exists but can't find when you need it most.

After experiencing this pain point repeatedly in my development work, I decided to build NeuroStash: an open-source knowledge management system designed to ingest, process, and intelligently extract information from document which can be self-hosted. But this isn't just another file storage solution—it's a production-ready system built with scalability, security, and developer experience in mind.

Project structure

NeuroStash
|- app/ # HTTP Endpoints and request/respones handler
|- core/ # Configuration and Database connection
|- dao/ # Data Access Objects - Database Operations
|- consumer/ # Asynchronous message processing
|- processor/ # Document processing pipeline
|- migrations # Database schema migrations

Database Schema Design: Multi-tenancy and Security First

The core of NeuroStash lies in its carefully designed database schema that balances performance, security, and flexibility

Multi-tenant User Management

class UserClient(Base, TimestampMixin):
    __tablename__ = "user_clients"

    id: Mapped[int] = mapped_column(Integer, Identity(), primary_key=True)
    email: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)
    role: Mapped[ClientRoleEnum] = mapped_column(
        SQLEnum(ClientRoleEnum, name="client_roles", create_type=False), 
        nullable=False
    )

    api_keys: Mapped[List["ApiKey"]] = relationship(
        back_populates="user_client", cascade="all, delete-orphan"
    )
    documents: Mapped[List["DocumentRegistry"]] = relationship(
        back_populates="user_client"
    )

Document Lifecycle Management

class DocumentRegistry(Base, TimestampMixin):
    __tablename__ = "documents_registry"

    user_id: Mapped[int] = mapped_column(ForeignKey("user_clients.id"))
    file_name: Mapped[str] = mapped_column(String(100), nullable=False)
    object_key: Mapped[str] = mapped_column(String(150), nullable=False)

    lock_status: Mapped[bool] = mapped_column(Boolean, nullable=False)
    op_status: Mapped[OperationStatusEnum] = mapped_column(
        SQLEnum(OperationStatusEnum), 
        nullable=False,
        server_default=OperationStatusEnum.PENDING.value
    )

    __table_args__ = (
        UniqueConstraint("user_id", "file_name", name="idx_unique_filename"),
        Index("idx_file_registry_user_id", "user_id", "lock_status", "op_status"),
    )

The lock_status field prevents concurrent processing of the same document, while op_status provides clear visibility into processing state.
The compound index on (user_id, lock_status, op_status) enables efficient queries for finding document ready for processing

Encryption Key Management

class EncryptionKey(Base, TimestampMixin):
    __tablename__ = "encryption_keys"

    symmetric_key: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
    is_active: Mapped[bool] = mapped_column(Boolean, server_default=text("false"))
    expired_at: Mapped[Optional[datetime]] = mapped_column(TIMESTAMP(timezone=True))

    __table_args__ = (
        Index("idx_encryption_keys_active", "id", postgresql_where=Column("is_active")),
    )

The partial index on active keys ensures lightning-fast lookups for current encryption keys.

Knowledge Base Organization

class KnowledgeBaseDocument(Base, TimestampMixin):
    knowledge_base_id: Mapped[int] = mapped_column(ForeignKey("knowledge_bases.id"))
    document_id: Mapped[int] = mapped_column(ForeignKey("documents_registry.id"))
    status: Mapped[OperationStatusEnum] = mapped_column(
        SQLEnum(OperationStatusEnum),
        server_default=OperationStatusEnum.PENDING.value
    )

This design allows document to belong to multiple knowledge bases while tracking the processing status for each relationship independently — crucial for systems that need to rebuild indexes or reprocess documents

Configuration Management: Production Ready from the Start

Professional applications require robust configuration management and NeuroStash uses pydantic settings for type-safe, environment-aware configurations

class Settings(BaseSettings):
    model_config = SettingsConfigDict(
        env_file=os.path.join(os.path.dirname(__file__), "..", "..", ".env"),
        env_file_encoding="utf-8",
        extra="ignore",
    )

    POSTGRES_SERVER: str
    POSTGRES_PORT: int
    POSTGRES_USER: str
    POSTGRES_PASSWORD: str
    POSTGRES_DB: str

    @computed_field
    @property
    def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn:
        return MultiHostUrl.build(
            scheme="postgresql+psycopg",
            username=self.POSTGRES_USER,
            password=self.POSTGRES_PASSWORD,
            host=self.POSTGRES_SERVER,
            port=self.POSTGRES_PORT,
            path=self.POSTGRES_DB,
        )

Database Migrations

Database migrations are often an afterthought, but they’re critical for production systems.
NeuroStash uses Alembic with several key Optimizations

Explicit Enum Control Strategy

I learned this lesson the hard way: Alembic’s auto-generated enum handling has critical flaw during migration downgrades. When you downgrade a migration, the enums aren’t properly deleted, causing “enum already exists” errors when you try to run the migration again.

client_role_enum = postgresql.ENUM(
    "USER", "ADMIN", name="client_roles", create_type=False
)
operation_status_enum = postgresql.ENUM(
    "PENDING", "SUCCESS", "FAILED", name="operation_status", create_type=False
)

def upgrade() -> None:
    # Explicitly create enums with checkfirst=True
    client_role_enum.create(op.get_bind(), checkfirst=True)
    operation_status_enum.create(op.get_bind(), checkfirst=True)

    # Then create tables that reference them
    op.create_table("user_clients", ...)

def downgrade() -> None:
    # Drop tables first
    op.drop_table("user_clients")

    # Then explicitly drop enums
    operation_status_enum.drop(op.get_bind(), checkfirst=True)
    client_role_enum.drop(op.get_bind(), checkfirst=True)

When testing migrations locally, you can reliably downgrade and re-run without manual database cleanup
Production rollbacks work cleanly without manual intervention
checkfirst=True is very crucial — it prevent errors if the enum already exists during the upgrade or doesn’t exists during downgrade.

Migration Environment Setup

The env.py configuration demonstrates production-ready practices

db_url = str(settings.SQLALCHEMY_DATABASE_URI)
config.set_main_option("sqlalchemy.url", db_url)

Instead of passing the database connection url we managed through configuration and env.py

Database Connection Management

The database connection setup in db.py includes several production-critical configuration

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from app.core.config import settings

engine = create_engine(
    url=str(settings.SQLALCHEMY_DATABASE_URI),
    pool_pre_ping=True,
    connect_args={"options": "-c timezone=utc"},
)

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

pool_pre_ping=True: This is crucial for long-running applications. SQLAlchemy will test connections before use, automatically handling dropped connections that can occur in cloud environments or during database maintenance windows. Without this, you'd get cryptic connection errors in production.
timezone=utc: Forces all database connections to use UTC, preventing timezone-related bugs when the application runs in different regions or when daylight saving time changes occur. This is especially important for multi-tenant applications where users might be in different timezones.
Session Configuration: autocommit=False and autoflush=False provide explicit control over transaction boundaries, which is essential for maintaining data consistency in complex operations involving multiple tables.
autocommit determines each SQL statement is automatically committed to the database immediately after execution, rather than waiting for explicitly commit() call.
autoflush controls whether pending changes in session are automatically sent to the database befor executing queries

Database Initialization Strategy

The pre_start.py script implements a robust database readiness check

@retry(
    stop=stop_after_attempt(max_tries),
    wait=wait_fixed(wait_seconds),
    before=before_log(logger, logging.DEBUG),
    after=after_log(logger, logging.INFO),
)
def init(db_engine: Engine) -> None:
    try:
        with SqlSession(db_engine) as session:
            session.execute(text("SELECT 1"))
        logger.info("database connection successfully")
    except Exception as e:
        logger.error(f"database connection/query failed: {e}")
        raise

Current Status

Document Processing pipeline
Vector search integration
API Authentication

So much more…

Takeaways

Building NeuroStash is journey which is ongoing and has taught me plenty of things and yet more to come.

I learnt that you always start with strong foundation with proper configuration and database design and thinking long term, envisioning the product will pay dividends later

Building encryption, authentication, multi-tenancy solution from beginning is easier than retrofitting

Explicit control is better for production-systems

Check out the code: Code

Connect with me: X